{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. 3D Image Registration with Masks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Registration"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Often, some content in the images may not correspond. For example, there may be background content or noisy areas. A mask defined on the fixed and/or moving image can be used to exclude these regions at a pixel level from the similarity metric computations. This improves the robustness of registration.\n",
    "\n",
    "A mask is a binary image with 1 meaning that a pixel is include in elastix' computation and a 0 meaning it's not.\n",
    "\n",
    "For more information, see Section 5.4, \"Masks\" of the [Elastix Manual](https://elastix.lumc.nl/download/elastix-5.0.1-manual.pdf)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import itk\n",
    "from itkwidgets import compare, checkerboard, view"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The function calls in the 3D case to import and register the images is similar to the 2D case. Masks, usually binary images, are import with the itk library similar to the images.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import Images\n",
    "fixed_image = itk.imread('data/CT_3D_lung_fixed.mha', itk.F)\n",
    "moving_image = itk.imread('data/CT_3D_lung_moving.mha', itk.F)\n",
    "\n",
    "# Import Custom Parameter Map\n",
    "parameter_object = itk.ParameterObject.New()\n",
    "parameter_object.AddParameterFile('data/parameters.3D.NC.affine.ASGD.001.txt')\n",
    "\n",
    "# \"WriteResultImage\" needs to be set to \"true\" so that the image is resampled at the end of the registration\n",
    "# and the result_image is populated properly\n",
    "parameter_object.SetParameter(0, \"WriteResultImage\", \"true\")\n",
    "\n",
    "# Import Mask Images\n",
    "fixed_mask = itk.imread('data/CT_3D_lung_fixed_mask.mha', itk.UC)\n",
    "moving_mask = itk.imread('data/CT_3D_lung_moving_mask.mha', itk.UC)\n",
    "\n",
    "# Or Optionally Create Masks from scratch\n",
    "\n",
    "# MaskImageType = itk.Image[itk.UC, 2]\n",
    "# fixed_mask = itk.binary_threshold_image_filter(fixed,\n",
    "#                                               lower_threshold=80.0,\n",
    "#                                               inside_value=1,\n",
    "#                                               ttype=(type(fixed), MaskImageType))\n",
    "# moving_mask = itk.binary_threshold_image_filter(moving,\n",
    "#                                                lower_threshold=80.0,\n",
    "#                                                inside_value=1,\n",
    "#                                                ttype=(type(moving), MaskImageType))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Input Visualization\n",
    "The images and their masks can be visualized with the itkwidget's view function. This can be useful to visually inspect the quality of the masks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9e0ed89550c043a29bef90625e1cb58e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Viewer(geometries=[], gradient_opacity=0.22, interpolation=False, point_sets=[], rendered_image=<itk.itkImageP…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "view(fixed_image, label_image=fixed_mask)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3e093036b3be428389dcebe081bb587b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Viewer(geometries=[], gradient_opacity=0.22, interpolation=False, point_sets=[], rendered_image=<itk.itkImageP…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "view(moving_image, label_image = moving_mask)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Registration"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Registration can either be done in one line with the registration function..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Call registration function\n",
    "result_image, result_transform_parameters = itk.elastix_registration_method(\n",
    "    fixed_image, moving_image,\n",
    "    parameter_object=parameter_object,\n",
    "    fixed_mask=fixed_mask, moving_mask=moving_mask,\n",
    "    log_to_console=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ".. or by initiating an elastix image filter object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load Elastix Image Filter Object\n",
    "# Fixed and moving image should be given to the Elastix method to ensure that\n",
    "# the correct 3D class is initialized.\n",
    "elastix_object = itk.ElastixRegistrationMethod.New(fixed_image, moving_image)\n",
    "elastix_object.SetFixedMask(fixed_mask)\n",
    "elastix_object.SetMovingMask(moving_mask)\n",
    "elastix_object.SetParameterObject(parameter_object)\n",
    "\n",
    "# Set additional options\n",
    "elastix_object.SetLogToConsole(False)\n",
    "\n",
    "# Update filter object (required)\n",
    "elastix_object.UpdateLargestPossibleRegion()\n",
    "\n",
    "# Results of Registration\n",
    "result_image = elastix_object.GetOutput()\n",
    "result_transform_parameters = elastix_object.GetTransformParameterObject()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Output Visualization\n",
    "The results of the 3D image registration can also be visualized with widgets from the itkwidget library such as the checkerboard and compare widgets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "94d279f136cc48b3983cc60458776d58",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(Viewer(annotations=False, interpolation=False, rendered_image=<itk.itkImagePython.itkImageF3; p…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "checkerboard(fixed_image, result_image,pattern=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b68308d4d978479789642fe885a49896",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "AppLayout(children=(HBox(children=(Label(value='Link:'), Checkbox(value=True, description='cmap'), Checkbox(va…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "compare(fixed_image, result_image, label_image= [fixed_image, result_image],link_cmap=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}