{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Rigid registration\n",
    "\n",
    "This example should show you how to perform rigid registration between two SIRF images. It is recommended that you check the [Geometry notebooks](../Geometry) first.\n",
    "\n",
    "SIRF's registration/resampling functionality is provided by wrapping and extending the [NiftyReg](http://cmictig.cs.ucl.ac.uk/wiki/index.php/NiftyReg) code base. Rigid and affine registrations are performed using NiftyReg's symmetric `aladin` algorithm, whereas non-rigid registrations use the symmetric `f3d` algorithm.\n",
    "\n",
    "Although the example of rigid registration is given here, it is trivial to modify it for affine registrations and only slightly trickier to extend to non-rigid registration.\n",
    "\n",
    "The images to be registered in this example are `test.nii.gz` and `test2.nii.gz`, which are two T1-weighted MR brain scans taken one year apart.\n",
    "\n",
    "N.B.: Registration packages use different names for the sets of images they are registering. In NiftyReg (and therefore SIRF), the floating image is moved to match the reference image. In other packages, the floating=moving and reference=fixed.\n",
    "\n",
    "N.B.: If you have MATLAB, you can install SPM. `SIRF` will then be able to wrap the `spm_realign` function as well (even from Python!). This is not illustrated here, but see the `sirf.Reg.SPMRegistration` class documentation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Copyright\n",
    "\n",
    "Author: Richard Brown  \n",
    "First version: 3rd April 2019  \n",
    "Second version: 8th June 2020  \n",
    "\n",
    "CCP PETMR Synergistic Image Reconstruction Framework (SIRF)  \n",
    "Copyright 2019-2021 University College London.\n",
    "\n",
    "This is software developed for the Collaborative Computational Project in Positron Emission Tomography and Magnetic Resonance imaging (http://www.ccppetmr.ac.uk/).\n",
    "\n",
    "SPDX-License-Identifier: Apache-2.0\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initial set up"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#%% make sure figures appears inline and animations works\n",
    "%matplotlib widget\n",
    "\n",
    "# Setup the working directory for the notebook\n",
    "import notebook_setup\n",
    "from sirf_exercises import cd_to_working_dir\n",
    "cd_to_working_dir('Reg', 'Registration')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Standard stuff\n",
    "import os\n",
    "import numpy\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# SIRF stuff\n",
    "from sirf.Utilities import examples_data_path\n",
    "import sirf.Reg as Reg\n",
    "examples_path = examples_data_path('Registration')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#%% First define some handy function definitions\n",
    "# To make subsequent code cleaner, we have a few functions here. You can\n",
    "# ignore them when you first see this demo.\n",
    "# They have (minimal) documentation using Python docstrings such that you \n",
    "# can do for instance \"help(imshow)\"\n",
    "#\n",
    "# First a function to display an image\n",
    "\n",
    "def imshow(image, title=''):\n",
    "    \"\"\"Display an image with a colourbar, returning the plot handle. \n",
    "\n",
    "    Arguments:\n",
    "    image -- a 2D array of numbers\n",
    "    limits -- colourscale limits as [min,max]. An empty [] uses the full range\n",
    "    title -- a string for the title of the plot (default \"\")\n",
    "    \"\"\"\n",
    "    plt.title(title)\n",
    "    bitmap=plt.imshow(image)\n",
    "    limits=[numpy.nanmin(image),numpy.nanmax(image)]\n",
    "\n",
    "    plt.clim(limits[0], limits[1])\n",
    "    plt.colorbar(shrink=.6)\n",
    "    plt.axis('off')\n",
    "    return bitmap"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Engine for loading images\n",
    "\n",
    "By default this example uses `sirf.Reg` as the engine to open the images, which handles NIfTI images. \n",
    "\n",
    "You might want to register different types of images - perhaps your floating image is a STIR interfile? If so, change the second line such that it reads `import sirf.STIR as eng_flo`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# In this example, we will use sirf.Reg as the engine to open our images\n",
    "import sirf.Reg as eng_ref\n",
    "import sirf.Reg as eng_flo"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Open and display the images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Open the images\n",
    "ref_file = os.path.join(examples_path, \"test.nii.gz\")\n",
    "flo_file = os.path.join(examples_path, \"test2.nii.gz\")\n",
    "ref = eng_ref.ImageData(ref_file)\n",
    "flo = eng_flo.ImageData(flo_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Display the images\n",
    "ref_slice = ref.get_dimensions()[1] // 2\n",
    "flo_slice = flo.get_dimensions()[1] // 2\n",
    "plt.subplot(1,2,1)\n",
    "imshow(ref.as_array()[ref_slice,:,:], 'Reference image, slice: %i' % ref_slice)\n",
    "plt.subplot(1,2,2)\n",
    "imshow(flo.as_array()[flo_slice,:,:], 'Floating image, slice: %i' % flo_slice);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A little word about image orientation\n",
    "\n",
    "An image will needs information in order to be able to place it in the real world. This information encodes things such as voxel sizes, orientation and offset or origin (where the image starts). The geometrical information for any SIRF image can be extracted with `get_geometrical_info`.\n",
    "\n",
    "This is discussed in detail in the [Geometry notebooks](../Geometry).\n",
    "\n",
    "Let's have a look at the direction matrix and offset for our two example images:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Ref\n",
    "ref_geom_info = ref.get_geometrical_info()\n",
    "ref_direction_matrix = ref_geom_info.get_direction_matrix()\n",
    "print(\"printing ref direction matrix\\n\", ref_direction_matrix)\n",
    "ref_offset = ref_geom_info.get_offset()\n",
    "print(\"printing ref offset\\n\", ref_offset)\n",
    "# Flo\n",
    "flo_geom_info = flo.get_geometrical_info()\n",
    "flo_direction_matrix = flo_geom_info.get_direction_matrix()\n",
    "print(\"\\nprinting flo direction matrix\\n\", flo_direction_matrix)\n",
    "flo_offset = flo_geom_info.get_offset()\n",
    "print(\"printing flo offset\\n\", flo_offset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This means that our `imshow` above is misleading. It uses `as_array`, which returns the underlying data \"as is\", without any regard for the image orientation. Luckily, the registration and resampling tools understand this and will take it automatically into account."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Hopefully you can convince yourself that, in spite of the misleading orientation of the first `imshow`, the images are actually facing the same direction. \n",
    "\n",
    "Now have a look at our offsets that we printed further up. These are quite different. We would therefore expect the rigid registration between these two images would result in little rotation, but quite a lot of translation. \n",
    "\n",
    "Let's have a look!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Back to registration!\n",
    "\n",
    "Ok, let's get back to the registration. Start by setting up the registration object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set to NiftyF3dSym for non-rigid\n",
    "algo = Reg.NiftyAladinSym()\n",
    "\n",
    "# Set images\n",
    "algo.set_reference_image(ref)\n",
    "algo.set_floating_image(flo)\n",
    "\n",
    "# What else can we do?\n",
    "help(algo)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting parameters\n",
    "\n",
    "From the help above, it looks like we can set registration parameters both via a file and directly. Let's try both."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setting via parameter file\n",
    "par_file = os.path.join(examples_path, \"paramFiles\", \"niftyreg_aladin.par\")\n",
    "algo.set_parameter_file(par_file)\n",
    "\n",
    "algo.set_parameter('SetPerformRigid','1')\n",
    "algo.set_parameter('SetPerformAffine','0')\n",
    "#algo.set_parameter('SetWarpedPaddingValue','0')  # NaN by default, uncomment to set to 0."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Register!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "algo.process()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Show the registered image\n",
    "\n",
    "The registered image will be the same size as the reference image, so we can use `ref_slice`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output = algo.get_output()\n",
    "output_arr = output.as_array()\n",
    "plt.figure()\n",
    "imshow(output_arr[ref_slice,:,:], 'Registered image, slice: %i' % ref_slice);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Show the transformation matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "numpy.set_printoptions(precision=3,suppress=True)\n",
    "TM = algo.get_transformation_matrix_forward()\n",
    "print(TM.as_array())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So as we predicted earlier, there is little rotation (close to identity in the top 3x3 of the matrix), and there is large translation (final column)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Displacement field image\n",
    "The displacement field image maps the voxels of the floating image to the final warped image. They are particularly interesting for non-rigid registrations, where transformation matrices no longer exist.\n",
    "\n",
    "Each voxel of the displacement field image contains an (x,y,z) coordinate, so the resulting image is 4D.\n",
    "\n",
    "(As a small technicality, the NIfTI format stores the time component in the 4th dimension, and the displacement coordinates in the 5th dimension. Therefore the displacement field image is actually a 5D image with a singleton in the 4th dimension.)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "***WARNING***: the displacement vectors from `sirf.Reg` are currently (SIRF 3.1) in RAS coordinates as used by NIfTI."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "displacement = algo.get_displacement_field_forward()\n",
    "\n",
    "plt.figure()\n",
    "for i in range(3):\n",
    "    plt.subplot(3, 1, 1 + i)\n",
    "    imshow(displacement.as_array()[ref_slice,:,:,0,i],\n",
    "           'Displacement field %s-direction, slice: %i' % (\"xyz\"[i], ref_slice))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Resampling\n",
    "\n",
    "Imagine that you got the motion information from another source and wanted to resample an image with it. Easy!\n",
    "\n",
    "We just need to create a resampler and set our image to resample as a floating image. As image domains can be different (i.e. images from different modalities, images with different field of view, etc), we also need to give the resampler a \"template\" image to define the resampled space. We can do this by setting the reference image. If we want to resample the image to the same domain as the original, you can set the floating image and the reference image to be the same.\n",
    "\n",
    "We would also need to define the type of transformation. We can use transformation matrices here, but it can also be a displacement of a deformation field, if your application requires it. \n",
    "\n",
    "Finally, we can also define the type of interpolation to use. In this example, we use linear with `set_interpolation_type_to_linear()`, but we can also use `_nearest_neighbour()`, `_cubic_spline()` and `_sinc()` instead of `_linear()`. Or just set it with `set_interpolation_type(int)`, with their respective order of interpolation.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tm = algo.get_transformation_matrix_forward()\n",
    "# get the resampler!\n",
    "resampler = Reg.NiftyResample()\n",
    "# Make sure we know what the resampled image domain looks like (this can be the same as the image to resample)\n",
    "resampler.set_reference_image(ref)\n",
    "# Set image to resample\n",
    "resampler.set_floating_image(flo)\n",
    "# Add the desired transformation to apply\n",
    "resampler.add_transformation(tm)\n",
    "resampler.set_padding_value(0)\n",
    "# Use linear interpolation\n",
    "resampler.set_interpolation_type_to_linear()\n",
    "# Go!\n",
    "resampler.process()\n",
    "\n",
    "plt.figure()\n",
    "imshow(resampler.get_output().as_array()[ref_slice,:,:], 'Resampled image, slice: %i' % ref_slice);"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython"
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}