{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Crossing-preserving contextual enhancement\n\nThis demo presents an example of crossing-preserving contextual enhancement of\nFOD/ODF fields [Meesters2016]_, implementing the contextual PDE framework\nof [Portegies2015a]_ for processing HARDI data. The aim is to enhance the\nalignment of elongated structures in the data such that crossing/junctions are\nmaintained while reducing noise and small incoherent structures. This is\nachieved via a hypo-elliptic 2nd order PDE in the domain of coupled positions\nand orientations $\\mathbb{R}^3 \rtimes S^2$. This domain carries a\nnon-flat geometrical differential structure that allows including a notion of\nalignment between neighboring points.\n\nLet $({\bf y},{\bf n}) \\in \\mathbb{R}^3\rtimes S^2$ where\n${\bf y} \\in \\mathbb{R}^{3}$ denotes the spatial part, and\n${\bf n} \\in S^2$ the angular part.\nLet $W:\\mathbb{R}^3\rtimes S^2       imes \\mathbb{R}^{+}     o \\mathbb{R}$ be\nthe function representing the evolution of FOD/ODF field. Then, the contextual\nPDE with evolution time $t\\geq 0$ is given by:\n\n\\begin{align}\begin{cases}\n  \frac{\\partial}{\\partial t} W({\bf y},{\bf n},t) &= ((D^{33}({\bf n} \\cdot\\end{align}\nabla)^2 + D^{44} \\Delta_{S^2})W)({\bf y},{\bf n},t)\n  \\ W({\bf y},{\bf n},0) &= U({\bf y},{\bf n})\n  \\end{cases},\n\nwhere:\n\n* $D^{33}>0$ is  the coefficient for the spatial smoothing (which goes only in the direction of $n$);\n\n* $D^{44}>0$ is the coefficient for the angular smoothing (here $\\Delta_{S^2}$ denotes the Laplace-Beltrami operator on the sphere $S^2$);\n\n* $U:\\mathbb{R}^3\rtimes S^2       o \\mathbb{R}$ is the initial condition given by the noisy FOD/ODF\u2019s field.\n\nThis equation is solved via a shift-twist convolution (denoted by $\u0007st_{\\mathbb{R}^3\rtimes S^2}$) with its corresponding kernel $P_t:\\mathbb{R}^3\rtimes S^2       o \\mathbb{R}^+$:\n\n\\begin{align}W({\bf y},{\bf n},t) = (P_t \u0007st_{\\mathbb{R}^3 \rtimes S^2} U)({\bf y},{\bf n})\n  = \\int_{\\mathbb{R}^3} \\int_{S^2} P_t (R^T_{{\bf n}^\\prime}({\bf y}-{\bf y}^\\prime),\n   R^T_{{\bf n}^\\prime} {\bf n} ) U({\bf y}^\\prime, {\bf n}^\\prime)\\end{align}\n\nHere, $R_{\bf n}$ is any 3D rotation that maps the vector $(0,0,1)$\nonto ${\bf n}$.\n\nNote that the shift-twist convolution differs from a Euclidean convolution and\ntakes into account the non-flat structure of the space $\\mathbb{R}^3\rtimes S^2$.\n\nThe kernel $P_t$ has a stochastic interpretation [DuitsAndFranken2011]_.\nIt can be seen as the limiting distribution obtained by accumulating random\nwalks of particles in the position/orientation domain, where in each step the\nparticles can (randomly) move forward/backward along their current orientation,\nand (randomly) change their orientation. This is an extension to the 3D case of\nthe process for contour enhancement of 2D images.\n\n.. figure:: _static/stochastic_process.jpg\n   :scale: 60 %\n   :align: center\n\n   The random motion of particles (a) and its corresponding probability map\n   (b) in 2D. The 3D kernel is shown on the right. Adapted from\n   [Portegies2015a]_.\n\nIn practice, as the exact analytical formulas for the kernel $P_t$\nare unknown, we use the approximation given in [Portegies2015b]_.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import numpy as np\nfrom dipy.core.gradients import gradient_table\nfrom dipy.data import get_fnames, default_sphere\nfrom dipy.io.image import load_nifti_data\nfrom dipy.io.gradients import read_bvals_bvecs\nfrom dipy.sims.voxel import add_noise\nfrom dipy.reconst.csdeconv import odf_sh_to_sharp"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The enhancement is evaluated on the Stanford HARDI dataset\n(150 orientations, b=2000 $s/mm^2$) where Rician noise is added. Constrained\nspherical deconvolution is used to model the fiber orientations.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Read data\nhardi_fname, hardi_bval_fname, hardi_bvec_fname = get_fnames('stanford_hardi')\ndata = load_nifti_data(hardi_fname)\nbvals, bvecs = read_bvals_bvecs(hardi_bval_fname, hardi_bvec_fname)\ngtab = gradient_table(bvals, bvecs)\n\n# Add Rician noise\nfrom dipy.segment.mask import median_otsu\nb0_slice = data[:, :, :, 1]\nb0_mask, mask = median_otsu(b0_slice)\nnp.random.seed(1)\ndata_noisy = add_noise(data, 10.0, np.mean(b0_slice[mask]), noise_type='rician')\n\n# Select a small part of it.\npadding = 3  # Include a larger region to avoid boundary effects\ndata_small = data[25-padding:40+padding, 65-padding:80+padding, 35:42]\ndata_noisy_small = data_noisy[25-padding:40+padding,\n                              65-padding:80+padding,\n                              35:42]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Enables/disables interactive visualization\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "interactive = False"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Fit an initial model to the data, in this case Constrained Spherical\nDeconvolution is used.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Perform CSD on the original data\nfrom dipy.reconst.csdeconv import auto_response_ssst\nfrom dipy.reconst.csdeconv import ConstrainedSphericalDeconvModel\nresponse, ratio = auto_response_ssst(gtab, data, roi_radii=10, fa_thr=0.7)\ncsd_model_orig = ConstrainedSphericalDeconvModel(gtab, response)\ncsd_fit_orig = csd_model_orig.fit(data_small)\ncsd_shm_orig = csd_fit_orig.shm_coeff\n\n# Perform CSD on the original data + noise\nresponse, ratio = auto_response_ssst(gtab, data_noisy, roi_radii=10, fa_thr=0.7)\ncsd_model_noisy = ConstrainedSphericalDeconvModel(gtab, response)\ncsd_fit_noisy = csd_model_noisy.fit(data_noisy_small)\ncsd_shm_noisy = csd_fit_noisy.shm_coeff"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Inspired by [Rodrigues2010]_, a lookup-table is created, containing\nrotated versions of the kernel $P_t$ sampled over a discrete set of\norientations. In order to ensure rotationally invariant processing, the\ndiscrete orientations are required to be equally distributed over a sphere.\nBy default, a sphere with 100 directions is used.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from dipy.denoise.enhancement_kernel import EnhancementKernel\nfrom dipy.denoise.shift_twist_convolution import convolve\n\n# Create lookup table\nD33 = 1.0\nD44 = 0.02\nt = 1\nk = EnhancementKernel(D33, D44, t)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Visualize the kernel\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from dipy.viz import window, actor\nfrom dipy.reconst.shm import sf_to_sh, sh_to_sf\nscene = window.Scene()\n\n# convolve kernel with delta spike\nspike = np.zeros((7, 7, 7, k.get_orientations().shape[0]), dtype=np.float64)\nspike[3, 3, 3, 0] = 1\nspike_shm_conv = convolve(sf_to_sh(spike, k.get_sphere(), sh_order=8), k,\n                          sh_order=8, test_mode=True)\n\nspike_sf_conv = sh_to_sf(spike_shm_conv, default_sphere, sh_order=8)\nmodel_kernel = actor.odf_slicer(spike_sf_conv * 6,\n                                sphere=default_sphere,\n                                norm=False,\n                                scale=0.4)\nmodel_kernel.display(x=3)\nscene.add(model_kernel)\nscene.set_camera(position=(30, 0, 0), focal_point=(0, 0, 0), view_up=(0, 0, 1))\nwindow.record(scene, out_path='kernel.png', size=(900, 900))\nif interactive:\n    window.show(scene)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        ".. figure:: kernel.png\n   :align: center\n\n   Visualization of the contour enhancement kernel.\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Shift-twist convolution is applied on the noisy data\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Perform convolution\ncsd_shm_enh = convolve(csd_shm_noisy, k, sh_order=8)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The Sharpening Deconvolution Transform is applied to sharpen the ODF field.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Sharpen via the Sharpening Deconvolution Transform\n\ncsd_shm_enh_sharp = odf_sh_to_sharp(csd_shm_enh, default_sphere, sh_order=8,\n                                    lambda_=0.1)\n\n# Convert raw and enhanced data to discrete form\ncsd_sf_orig = sh_to_sf(csd_shm_orig, default_sphere, sh_order=8)\ncsd_sf_noisy = sh_to_sf(csd_shm_noisy, default_sphere, sh_order=8)\ncsd_sf_enh = sh_to_sf(csd_shm_enh, default_sphere, sh_order=8)\ncsd_sf_enh_sharp = sh_to_sf(csd_shm_enh_sharp, default_sphere, sh_order=8)\n\n# Normalize the sharpened ODFs\ncsd_sf_enh_sharp = csd_sf_enh_sharp * np.amax(csd_sf_orig)/np.amax(csd_sf_enh_sharp) * 1.25"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The end results are visualized. It can be observed that the end result after\ndiffusion and sharpening is closer to the original noiseless dataset.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "scene = window.Scene()\n\n# original ODF field\nfodf_spheres_org = actor.odf_slicer(csd_sf_orig,\n                                    sphere=default_sphere,\n                                    scale=0.4,\n                                    norm=False)\nfodf_spheres_org.display(z=3)\nfodf_spheres_org.SetPosition(0, 25, 0)\nscene.add(fodf_spheres_org)\n\n# ODF field with added noise\nfodf_spheres = actor.odf_slicer(csd_sf_noisy,\n                                sphere=default_sphere,\n                                scale=0.4,\n                                norm=False,)\nfodf_spheres.SetPosition(0, 0, 0)\nscene.add(fodf_spheres)\n\n# Enhancement of noisy ODF field\nfodf_spheres_enh = actor.odf_slicer(csd_sf_enh,\n                                    sphere=default_sphere,\n                                    scale=0.4,\n                                    norm=False)\nfodf_spheres_enh.SetPosition(25, 0, 0)\nscene.add(fodf_spheres_enh)\n\n# Additional sharpening\nfodf_spheres_enh_sharp = actor.odf_slicer(csd_sf_enh_sharp,\n                                          sphere=default_sphere,\n                                          scale=0.4,\n                                          norm=False)\nfodf_spheres_enh_sharp.SetPosition(25, 25, 0)\nscene.add(fodf_spheres_enh_sharp)\n\nwindow.record(scene, out_path='enhancements.png', size=(900, 900))\nif interactive:\n    window.show(scene)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        ".. figure:: enhancements.png\n   :align: center\n\n   The results after enhancements. Top-left: original noiseless data.\n   Bottom-left: original data with added Rician noise (SNR=10). Bottom-right:\n   After enhancement of noisy data. Top-right: After enhancement and sharpening\n   of noisy data.\n\n## References\n\n.. [Meesters2016] S. Meesters, G. Sanguinetti, E. Garyfallidis, J. Portegies,\n   R. Duits. (2016) Fast implementations of contextual PDE\u2019s for HARDI data\n   processing in DIPY. ISMRM 2016 conference.\n\n.. [Portegies2015a] J. Portegies, R. Fick, G. Sanguinetti, S. Meesters,\n   G.Girard, and R. Duits. (2015) Improving Fiber Alignment in HARDI by\n   Combining Contextual PDE flow with Constrained Spherical Deconvolution.\n   PLoS One.\n\n.. [Portegies2015b] J. Portegies, G. Sanguinetti, S. Meesters, and R. Duits.\n   (2015) New Approximation of a Scale Space Kernel on SE(3) and Applications\n   in Neuroimaging. Fifth International Conference on Scale Space and\n   Variational Methods in Computer Vision.\n\n.. [DuitsAndFranken2011] R. Duits and E. Franken (2011) Left-invariant\n   diffusions on the space of positions and orientations and their application\n   to crossing-preserving smoothing of HARDI images. International Journal of\n   Computer Vision, 92:231-264.\n\n.. [Rodrigues2010] P. Rodrigues, R. Duits, B. Romeny, A. Vilanova (2010).\n   Accelerated Diffusion Operators for Enhancing DW-MRI. Eurographics Workshop\n   on Visual Computing for Biology and Medicine. The Eurographics Association.\n\n\n"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "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.13"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}