{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Implementing Filtered Back Projection"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll use the classes of PyTomography to implement filtered back projection in SPECT."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from pytomography.projections import SPECTSystemMatrix\n",
    "from pytomography.metadata import ObjectMeta, ImageMeta\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import torch"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The two foundational tools of image reconstruction are\n",
    "\n",
    "1. Forward projection $\\sum_{i} c_{ij} a_i$\n",
    "2. Back projection $\\sum_{j} c_{ij} b_j$\n",
    "\n",
    "Let's discuss what these operators actually mean. First, let's define our quantities. $c_{ij}$ is known as the system matrix, and may include information involving attenuation and PSF correction. $a_i$ is an arbtriary object and $b_j$ is an arbitrary image. \n",
    "\n",
    "It's worth now discussing what the indices $i$ and $j$ actually mean. You might think: objects are three dimensional, shouldn't there be at least 3 indices when we're doing linear operations? Consider the following: because we are in a discrete space, any 3 dimensional object can be converted to a single (albeit very long) one dimensional object: a 128x128x128 3D matrix can be converted into a single 1D vector of length 2097152. That's how many voxels there are in object space: you can think of index $i$ as indexing a single voxel.\n",
    "\n",
    "The same can be said for an image. If we have 64 projections of matrix size 128x128, then that can be thought of as a single vector of length 1048576. That's also how many individual detector elements there are.\n",
    "\n",
    "So in forward projection $\\sum_{i} c_{ij} a_i$, the system matrix $c_{ij}$ maps the contribution from voxel $i$ to a detector element $j$. In back projection $\\sum_{j} c_{ij} b_j$, the system matrix $c_{ij}$ maps the intensity in detector element back to every possible voxel $i$ that could have contributed to it. In reality, however, not every voxel that could have contributed to detector element $j$ does so with equal intensity; it is for this reason that forward projection followed by back projection does not yield the original image. \n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's experiment with these operators. First we'll make a 3D rectangle in object space:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 128, 128, 132])"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = torch.linspace(-1,1,128)\n",
    "y = torch.linspace(-1,1,128)\n",
    "z = torch.linspace(-1,1,132)\n",
    "xv, yv, zv = torch.meshgrid([x,y,z], indexing='ij')\n",
    "object_truth = (xv>-0.2)*(xv<0.2)*(yv>-0.15)*(yv<0.15)*(zv>-0.1)*(zv<0.1)\n",
    "object_truth = object_truth.to(torch.float).unsqueeze(dim=0) # add batch dimension\n",
    "object_truth.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x7f5955a0d070>"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 7484 characters]",
      "text/plain": [
       "<Figure size 500x400 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.figure(figsize=(5,4))\n",
    "plt.pcolormesh(object_truth[0][:,:,64].T, cmap='Greys_r')\n",
    "plt.axis('off')\n",
    "plt.colorbar()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before we do any projections, we need to get corresponding metadata for our object. In this case, we'll assume the voxel sizes are 1cm $^3$. For our image space, we'll assume 60 projections are taken at angular spacing of 6 degrees."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "angles = np.arange(0,360.,6.)\n",
    "object_meta = ObjectMeta(dr=(1,1,1), shape=object_truth[0].shape)\n",
    "image_meta = ImageMeta(object_meta, angles=angles)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With this metadata, we can create our forward and back projection networks. We'll model no phenomenon for now."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "system_matrix = SPECTSystemMatrix(\n",
    "    obj2obj_transforms=[],\n",
    "    im2im_transforms=[],\n",
    "    object_meta=object_meta,\n",
    "    image_meta=image_meta)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now use the `forward` method of `system_matrix` to model $g=Hf$ (convert the object $f$ into an image $g$)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 60, 128, 132])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "image = system_matrix.forward(object_truth)\n",
    "image.shape"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can look at a projection at 60 degrees for example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 10848 characters]",
      "text/plain": [
       "<Figure size 1500x400 with 5 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, axes = plt.subplots(1,5,figsize=(15,4))\n",
    "for i, proj in enumerate([0,5,10,15,20]):\n",
    "    axes[i].pcolormesh(image[0][proj].T, cmap='Greys_r')\n",
    "    axes[i].set_title(f'Angle={image_meta.angles[proj]}')\n",
    "    axes[i].axis('off')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "At angles like 60 degrees, the cube is darkest in the center and lighter on the outside; this is like looking through a semi-transparent cube in real life, it's going to be darkest near the center if you're looking at it from an off angle."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also back project using the `backward` method of the `system_matrix` instance. Here we compute $H^T g$ and the normalization constant $H^T 1$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          ...,\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.]],\n",
       "\n",
       "         [[0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          ...,\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.]],\n",
       "\n",
       "         [[0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          ...,\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.]],\n",
       "\n",
       "         ...,\n",
       "\n",
       "         [[0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          ...,\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.]],\n",
       "\n",
       "         [[0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          ...,\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.]],\n",
       "\n",
       "         [[0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          ...,\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.]]]])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "object_bp, norm_bp = system_matrix.backward(image, return_norm_constant=True)\n",
    "object_bp"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But if we look at the new object:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "axial_slice = object_bp[0][:,:,64].T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x7f5921114fa0>"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 26396 characters]",
      "text/plain": [
       "<Figure size 500x400 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.figure(figsize=(5,4))\n",
    "plt.pcolormesh(axial_slice, cmap='Greys_r')\n",
    "plt.axis('off')\n",
    "plt.colorbar()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that it has been blurred, and that the values are much too large. How can we fix this?"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Example**: Filtered Back Projection. In this case the image estimate is given by\n",
    "\n",
    "$$\\hat{f}_i = \\frac{\\pi}{N_{\\text{proj}}} \\sum_i H_{ij} \\left( \\mathcal{F}^{-1}(|\\omega|\\mathcal{F}(g)) \\right)_i $$\n",
    "\n",
    "where the term in brackets involves applying a 1D convolution (in this case, multiplication in Fourier space with the Ramp filter) to the image along the $r$ axis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "freq_fft = torch.fft.fftfreq(image.shape[-2]) \n",
    "filter = torch.abs(freq_fft).reshape((-1,1))\n",
    "image_fft = torch.fft.fft(image, axis=-2)\n",
    "image_fft = image_fft* filter\n",
    "image_filtered = torch.fft.ifft(image_fft, axis=-2).real"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can back project and normalize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "object_fbp = system_matrix.backward(image_filtered) *np.pi / len(image_meta.angles)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x7f592105af40>"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 33928 characters]",
      "text/plain": [
       "<Figure size 500x400 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.figure(figsize=(5,4))\n",
    "plt.pcolormesh(object_fbp[0][:,:,64].T, cmap='Greys_r')\n",
    "plt.axis('off')\n",
    "plt.colorbar()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The cube is no longer blurred, but artifacts are present. Such artifacts are not present when using algorithms like OSEM for reconstruction."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "torch",
   "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.9.16"
  },
  "vscode": {
   "interpreter": {
    "hash": "324f5aa1485c6e81c0125c86ad1340c603dab08bc11babd97bd730cfafa2ec4a"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
