{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "36810dc9",
   "metadata": {},
   "source": [
    "# Transforms can be used to center image data or to align two images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "498ee278",
   "metadata": {},
   "outputs": [],
   "source": [
    "import shapeworks as sw\n",
    "import numpy as np\n",
    "import pyvista as pv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2f945d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "pv.set_jupyter_backend('static')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83a2df5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "DATA = \"../Data/\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60e765e6",
   "metadata": {},
   "source": [
    "## Center of mass transforms\n",
    "Centers the contents of a given image using the specified resampling method"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7bb1c011",
   "metadata": {},
   "outputs": [],
   "source": [
    "filename = DATA + \"seg.ellipsoid_14.nrrd\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c48dfc73",
   "metadata": {},
   "outputs": [],
   "source": [
    "img = sw.Image(filename)\n",
    "print(\"center: \", img.center())\n",
    "print(\"center of mass: \", img.centerOfMass())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d96e0ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "slices = sw.sw2vtkImage(img).slice_orthogonal(x=25,y=25,z=25)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58d41c9c",
   "metadata": {},
   "outputs": [],
   "source": [
    "p = pv.Plotter(shape=(1,3), border=False)\n",
    "\n",
    "p.subplot(0,0)\n",
    "p.add_text(\"before centering\", position='lower_left')\n",
    "p.add_mesh(slices[2], style='points', show_scalar_bar=False)\n",
    "p.show_grid()\n",
    "p.camera_position = 'xy'\n",
    "p.subplot(0,1)\n",
    "p.add_mesh(slices[0], style='points', show_scalar_bar=False)\n",
    "p.show_grid()\n",
    "p.camera_position = 'yz'\n",
    "p.subplot(0,2)\n",
    "p.add_mesh(slices[1], style='points', show_scalar_bar=False)\n",
    "p.show_grid()\n",
    "p.camera_position = 'zx'\n",
    "\n",
    "p.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ebf2c248",
   "metadata": {},
   "source": [
    "### create and apply transform, noting that center of mass is now roughly in the center of the image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43edcbd3",
   "metadata": {},
   "outputs": [],
   "source": [
    "xform = img.createCenterOfMassTransform()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90aeb203",
   "metadata": {},
   "outputs": [],
   "source": [
    "# now center of mass is roughly at the center\n",
    "img.applyTransform(xform)\n",
    "print(\"center of mass: \", img.centerOfMass())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5251f93",
   "metadata": {},
   "source": [
    "### visualization of different resampling methods"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5487219",
   "metadata": {},
   "outputs": [],
   "source": [
    "# linear interpolation\n",
    "img = sw.Image(filename)\n",
    "img.applyTransform(xform, sw.InterpolationType.Linear)\n",
    "lin = sw.sw2vtkImage(img)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9bdc44ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "# closest point resampling\n",
    "img = sw.Image(filename)\n",
    "img.applyTransform(xform, sw.InterpolationType.NearestNeighbor)\n",
    "nn = sw.sw2vtkImage(img)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0281c8e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "lin_slices = lin.slice_orthogonal(x=25,y=25,z=25)\n",
    "nn_slices = nn.slice_orthogonal(x=25,y=25,z=25)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9dc6af96",
   "metadata": {},
   "outputs": [],
   "source": [
    "p = pv.Plotter(shape=(2,3), border=False)\n",
    "\n",
    "p.subplot(0,0)\n",
    "p.add_text(\"linear\")\n",
    "p.add_mesh(lin_slices[2], style='points', show_scalar_bar=False)\n",
    "p.show_grid()\n",
    "p.camera_position = 'xy'\n",
    "p.subplot(0,1)\n",
    "p.add_mesh(lin_slices[0], style='points', show_scalar_bar=False)\n",
    "p.show_grid()\n",
    "p.camera_position = 'yz'\n",
    "p.subplot(0,2)\n",
    "p.add_mesh(lin_slices[1], style='points', show_scalar_bar=False)\n",
    "p.show_grid()\n",
    "p.camera_position = 'zx'\n",
    "\n",
    "p.subplot(1,0)\n",
    "p.add_text(\"nearest neighbor\")\n",
    "p.add_mesh(nn_slices[2], style='points', show_scalar_bar=False)\n",
    "p.show_grid()\n",
    "p.camera_position = 'xy'\n",
    "p.subplot(1,1)\n",
    "p.add_mesh(nn_slices[0], style='points', show_scalar_bar=False)\n",
    "p.show_grid()\n",
    "p.camera_position = 'yz'\n",
    "p.subplot(1,2)\n",
    "p.add_mesh(nn_slices[1], style='points', show_scalar_bar=False)\n",
    "p.show_grid()\n",
    "p.camera_position = 'zx'\n",
    "\n",
    "p.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3ffff9b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "4ddcd6a2",
   "metadata": {},
   "source": [
    "## Rigid registration transforms\n",
    "Aligns two images using their distance transforms by converting these to meshes, computing the alignment using iterative closest point (ICP), then resampling to move the center of mass of one to be aligned with the other\n",
    "\n",
    "NOTE: images used to compute the transform **must** be distance transforms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c398762f",
   "metadata": {},
   "outputs": [],
   "source": [
    "filename1 = DATA + \"seg.ellipsoid_17.nrrd\"\n",
    "filename2 = DATA + \"seg.ellipsoid_19.nrrd\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79d71f93",
   "metadata": {},
   "source": [
    "### compute distance transforms to be used to compute alignment of the segmentations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "163371e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "img1 = sw.Image(filename1)\n",
    "img2 = sw.Image(filename2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1e2a109",
   "metadata": {},
   "outputs": [],
   "source": [
    "img1_slices_orig = sw.sw2vtkImage(img1).slice_orthogonal(x=25,y=25,z=25)\n",
    "img2_slices_orig = sw.sw2vtkImage(img2).slice_orthogonal(x=25,y=25,z=25)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc9659cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "p = pv.Plotter(shape=(2,3), border=False)\n",
    "\n",
    "p.subplot(0,0)\n",
    "p.add_text(\"Image 1\")\n",
    "p.add_mesh(img1_slices_orig[2], style='points', show_scalar_bar=False)\n",
    "p.show_grid()\n",
    "p.camera_position = 'xy'\n",
    "p.subplot(0,1)\n",
    "p.add_mesh(img1_slices_orig[0], style='points', show_scalar_bar=False)\n",
    "p.show_grid()\n",
    "p.camera_position = 'yz'\n",
    "p.subplot(0,2)\n",
    "p.add_mesh(img1_slices_orig[1], style='points', show_scalar_bar=False)\n",
    "p.show_grid()\n",
    "p.camera_position = 'zx'\n",
    "\n",
    "p.subplot(1,0)\n",
    "p.add_text(\"Image 2\")\n",
    "p.add_mesh(img2_slices_orig[2], style='points', show_scalar_bar=False)\n",
    "p.show_grid()\n",
    "p.camera_position = 'xy'\n",
    "p.subplot(1,1)\n",
    "p.add_mesh(img2_slices_orig[0], style='points', show_scalar_bar=False)\n",
    "p.show_grid()\n",
    "p.camera_position = 'yz'\n",
    "p.subplot(1,2)\n",
    "p.add_mesh(img2_slices_orig[1], style='points', show_scalar_bar=False)\n",
    "p.show_grid()\n",
    "p.camera_position = 'zx'\n",
    "\n",
    "p.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da6202f3",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "360593af",
   "metadata": {},
   "source": [
    "### create distance transforms to be used to compute the images' alignment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26f56f9c",
   "metadata": {},
   "outputs": [],
   "source": [
    "dt1 = img1.computeDT(1.0)\n",
    "dt2 = img2.computeDT(1.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cdf542bb",
   "metadata": {},
   "source": [
    "### create image alignment transforms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b320950",
   "metadata": {},
   "outputs": [],
   "source": [
    "xform_1_to_2 = dt1.createRigidRegistrationTransform(dt2)\n",
    "xform_2_to_1 = dt2.createRigidRegistrationTransform(dt1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d33f77c6",
   "metadata": {},
   "source": [
    "### apply the transforms to the original images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a22214f",
   "metadata": {},
   "outputs": [],
   "source": [
    "img1 = sw.Image(filename1)\n",
    "img2 = sw.Image(filename2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e4b092d",
   "metadata": {},
   "outputs": [],
   "source": [
    "img1.applyTransform(xform_1_to_2)\n",
    "img2.applyTransform(xform_2_to_1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c8a5c87e",
   "metadata": {},
   "source": [
    "### visualize the results\n",
    "Notice how it even rotates the images into each others' centers of mass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a46e72f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "img1_slices = sw.sw2vtkImage(img1).slice_orthogonal(x=25,y=25,z=25)\n",
    "img2_slices = sw.sw2vtkImage(img2).slice_orthogonal(x=25,y=25,z=25)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51983016",
   "metadata": {},
   "outputs": [],
   "source": [
    "p = pv.Plotter(shape=(3,3), border=False)\n",
    "\n",
    "p.subplot(0,0)\n",
    "p.add_text(\"1 -> 2\")\n",
    "p.add_mesh(img1_slices[2], style='points', show_scalar_bar=False, opacity=0.75)\n",
    "p.add_mesh(img2_slices_orig[2], style='points', show_scalar_bar=False, opacity=0.25)\n",
    "p.show_grid()\n",
    "p.camera_position = 'xy'\n",
    "p.subplot(0,1)\n",
    "p.add_mesh(img1_slices[0], style='points', show_scalar_bar=False, opacity=0.75)\n",
    "p.add_mesh(img2_slices_orig[0], style='points', show_scalar_bar=False, opacity=0.25)\n",
    "p.show_grid()\n",
    "p.camera_position = 'yz'\n",
    "p.subplot(0,2)\n",
    "p.add_mesh(img1_slices[1], style='points', show_scalar_bar=False, opacity=0.75)\n",
    "p.add_mesh(img2_slices_orig[1], style='points', show_scalar_bar=False, opacity=0.25)\n",
    "p.show_grid()\n",
    "p.camera_position = 'zx'\n",
    "\n",
    "p.subplot(1,0)\n",
    "p.add_text(\"2 -> 1\")\n",
    "p.add_mesh(img1_slices_orig[2], style='points', show_scalar_bar=False, opacity=0.25)\n",
    "p.add_mesh(img2_slices[2], style='points', show_scalar_bar=False, opacity=0.75)\n",
    "p.show_grid()\n",
    "p.camera_position = 'xy'\n",
    "p.subplot(1,1)\n",
    "p.add_mesh(img1_slices_orig[0], style='points', show_scalar_bar=False, opacity=0.25)\n",
    "p.add_mesh(img2_slices[0], style='points', show_scalar_bar=False, opacity=0.75)\n",
    "p.show_grid()\n",
    "p.camera_position = 'yz'\n",
    "p.subplot(1,2)\n",
    "p.add_mesh(img1_slices_orig[1], style='points', show_scalar_bar=False, opacity=0.25)\n",
    "p.add_mesh(img2_slices[1], style='points', show_scalar_bar=False, opacity=0.75)\n",
    "p.show_grid()\n",
    "p.camera_position = 'zx'\n",
    "\n",
    "p.subplot(2,0)\n",
    "p.add_text(\"both\")\n",
    "p.add_mesh(img1_slices[2], style='points', show_scalar_bar=False, opacity=0.5)\n",
    "p.add_mesh(img2_slices[2], style='points', show_scalar_bar=False, opacity=0.5)\n",
    "p.show_grid()\n",
    "p.camera_position = 'xy'\n",
    "p.subplot(2,1)\n",
    "p.add_mesh(img1_slices[0], style='points', show_scalar_bar=False, opacity=0.5)\n",
    "p.add_mesh(img2_slices[0], style='points', show_scalar_bar=False, opacity=0.5)\n",
    "p.show_grid()\n",
    "p.camera_position = 'yz'\n",
    "p.subplot(2,2)\n",
    "p.add_mesh(img1_slices[1], style='points', show_scalar_bar=False, opacity=0.5)\n",
    "p.add_mesh(img2_slices[1], style='points', show_scalar_bar=False, opacity=0.5)\n",
    "p.show_grid()\n",
    "p.camera_position = 'zx'\n",
    "\n",
    "p.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "424b616f",
   "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.7.8"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
