Example of non-rigid registration.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Code to run basic non-rigid registration on images"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to run this code, you must first install simpleitk. To do this on a Diamond linux box, you must use a user install only as access is not granted for a system-wide install. This can also be done in a terminal from jupyterhub, but you will need to restart the server after installation.\n",
    "\n",
    "To install simpleitk, run: pip install simpleitk --user"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, import all necessary packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import hyperspy.api as hs\n",
    "import SimpleITK as sitk"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define function for non-rigid registration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def nonrigid(im_stack, demons_it = 20, filter_size = 5.0, max_it = 3):\n",
    "    \"Function to non-rigidly register a series of images.\n",
    "    \n",
    "    Input\n",
    "    -----\n",
    "    im_stack: 3D numpy array of data\n",
    "    demons_it: int\n",
    "    The number of iterations for the demons algorithm to perform.\n",
    "    filter_size: float\n",
    "    max_it: int\n",
    "    The number of iterations to apply the full non-rigid alignment algorithm.\n",
    "    \n",
    "    demons = sitk.DemonsRegistrationFilter()\n",
    "    demons.SetNumberOfIterations( demons_it )\n",
    "    # Standard deviation for Gaussian smoothing of displacement field\n",
    "    demons.SetStandardDeviations( filter_size )\n",
    "    \n",
    "    for j in range(max_it):\n",
    "        #Get stack average\n",
    "        av_im = sitk.GetImageFromArray(np.float32(sum(im_stack)/len(im_stack))) #Faster than numpy.mean for small arrays?\n",
    "        \n",
    "        out_stack = []\n",
    "        \n",
    "        for i in range(len(im_stack)):\n",
    "            \n",
    "            moving = sitk.GetImageFromArray(np.float32(im_stack[i]))\n",
    "            \n",
    "            displacementField = demons.Execute( av_im, moving )\n",
    "            \n",
    "            dispfield = sitk.GetArrayFromImage(displacementField)\n",
    "            \n",
    "            outTx = sitk.DisplacementFieldTransform( displacementField )\n",
    "            \n",
    "            resampler = sitk.ResampleImageFilter()\n",
    "            resampler.SetReferenceImage(av_im);\n",
    "            resampler.SetInterpolator(sitk.sitkLinear)\n",
    "            resampler.SetDefaultPixelValue(100)\n",
    "            resampler.SetTransform(outTx)\n",
    "            \n",
    "            out_stack.append(sitk.GetArrayFromImage(resampler.Execute(moving)))\n",
    "            \n",
    "        im_stack = out_stack\n",
    "        max_disp = np.max(dispfield)\n",
    "        print(max_disp)\n",
    "            \n",
    "        if max_disp < 0.3:\n",
    "            print(\"NRR stopped after \"+str(j)+\" iterations.\")\n",
    "            break\n",
    "    \n",
    "    return(out_stack)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Import multi-frame file using hyperspy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "filename = #Insert file location here\n",
    "images = hs.load(filename,stack=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "images.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to use the non-rigid registration code, the image series must first be rigidly aligned. Hyperspy has its own rigid registration function, used below, but this doesn't always work (particularly for highly periodic data). Rigid alignment can be done prior to using this notebook in any other software."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "images.align2D()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once rigidly aligned, the non-rigid resgitration can be run (see the above docstring for more information on parameters)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "images_nrr = nonrigid(images.data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The aligned series can then be summed and processed using atomap etc."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "image = images_nrr = hs.signals.Signal2D(images_nrr)\n",
    "image_sum = images.sum()\n",
    "image_sum.plot()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pspy-dev",
   "language": "python",
   "name": "pspy-dev"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}