{
 "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](http://elastix.isi.uu.nl/download/elastix-5.0.0-manual.pdf)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# First import is currently necessary to run ITKElastix on MacOs\n",
    "from itk import itkElastixRegistrationMethodPython\n",
    "import itk\n",
    "from itkwidgets import compare, checkerboard"
   ]
  },
  {
   "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": 8,
   "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",
    "# 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": [
    "Registration can either be done in one line with the registration function..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "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": 5,
   "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": [
    "### 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": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "02494a1888bb44dd9d4a0a8ac0f1dfbd",
       "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": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "db241d04acb44bffa664ba9c3c6af4ab",
       "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, link_cmap=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ElastixEnv",
   "language": "python",
   "name": "elastixenv"
  },
  "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}