{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Affine Registration with Masks\n\nThis example explains how to compute a transformation to register two 3D\nvolumes by maximization of their Mutual Information [Mattes03]_. The\noptimization strategy is similar to that implemented in ANTS [Avants11]_.\n\nWe will use masks to define which pixels are used in the Mutual Information.\nMasking can also be done for registration of 2D images rather than 3D volumes.\n\nMasking for registration is useful in a variety of circumstances. For example,\nin cardiac MRI, where it is usually used to specify a region of interest on a\n2D static image, e.g., the left ventricle in a short axis slice. This\nprioritizes registering the region of interest over other structures that move\nwith respect to the heart.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from os.path import join as pjoin\nimport numpy as np\nimport matplotlib.pyplot as plt\nfrom dipy.viz import regtools\nfrom dipy.data import fetch_stanford_hardi\nfrom dipy.io.image import load_nifti\nfrom dipy.align.imaffine import (AffineMap,\n                                 MutualInformationMetric,\n                                 AffineRegistration)\nfrom dipy.align.transforms import (TranslationTransform3D,\n                                   RigidTransform3D)\n\nfrom dipy.align import (affine_registration, translation,\n                        rigid, register_series)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's fetch a single b0 volume from the Stanford HARDI dataset.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "files, folder = fetch_stanford_hardi()\nstatic_data, static_affine, static_img = load_nifti(\n                                            pjoin(folder, 'HARDI150.nii.gz'),\n                                            return_img=True)\nstatic = np.squeeze(static_data)[..., 0]\n\n# pad array to help with this example\npad_by = 10\nstatic = np.pad(static, [(pad_by, pad_by), (pad_by, pad_by), (0, 0)],\n                mode='constant', constant_values=0)\n\nstatic_grid2world = static_affine"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's create a moving image by transforming the static image.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "affmat = np.eye(4)\naffmat[0, -1] = 4\naffmat[1, -1] = 12\ntheta = 0.1\nc, s = np.cos(theta), np.sin(theta)\naffmat[0:2, 0:2] = np.array([[c, -s], [s, c]])\naffine_map = AffineMap(affmat,\n                       static.shape, static_grid2world,\n                       static.shape, static_grid2world)\nmoving = affine_map.transform(static)\nmoving_affine = static_affine.copy()\nmoving_grid2world = static_grid2world.copy()\n\nregtools.overlay_slices(static, moving, None, 2,\n                        \"Static\", \"Moving\", \"deregistered.png\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        ".. figure:: deregistered.png\n   :align: center\n\n   Same images but misaligned.\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's make some registration settings.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "nbins = 32\nsampling_prop = None\nmetric = MutualInformationMetric(nbins, sampling_prop)\n\n# small number of iterations for this example\nlevel_iters = [100, 10]\nsigmas = [1.0, 0.0]\nfactors = [2, 1]\n\naffreg = AffineRegistration(metric=metric,\n                            level_iters=level_iters,\n                            sigmas=sigmas,\n                            factors=factors)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Now let's register these volumes together without any masking. For the purposes\nof this example, we will not provide an initial transformation based on centre\nof mass, but this would work fine with masks.\n\nNote that use of masks is not currently implemented for sparse sampling.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "transform = TranslationTransform3D()\ntransl = affreg.optimize(static, moving, transform, None,\n                         static_grid2world, moving_grid2world,\n                         starting_affine=None,\n                         static_mask=None, moving_mask=None)\ntransform = RigidTransform3D()\ntransl = affreg.optimize(static, moving, transform, None,\n                         static_grid2world, moving_grid2world,\n                         starting_affine=transl.affine,\n                         static_mask=None, moving_mask=None)\ntransformed = transl.transform(moving)\n\ntransformed = transl.transform(moving)\nregtools.overlay_slices(static, transformed, None, 2,\n                        \"Static\", \"Transformed\", \"transformed.png\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        ".. figure:: transformed.png\n   :align: center\n\n   Registration result.\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We can also use a pipeline to achieve the same thing. For convenience in this\ntutorial, we will define a function that runs the pipeline and makes a figure.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def reg_func(figname, static_mask=None, moving_mask=None):\n    \"\"\"Convenience function for registration using a pipeline.\n       Uses variables in global scope, except for static_mask and moving_mask.\n    \"\"\"\n\n    pipeline = [translation, rigid]\n\n    xformed_img, reg_affine = affine_registration(\n        moving,\n        static,\n        moving_affine=moving_affine,\n        static_affine=static_affine,\n        nbins=32,\n        metric='MI',\n        pipeline=pipeline,\n        level_iters=level_iters,\n        sigmas=sigmas,\n        factors=factors,\n        static_mask=static_mask,\n        moving_mask=moving_mask)\n\n    regtools.overlay_slices(static, xformed_img, None, 2,\n                            \"Static\", \"Transformed\", figname)\n\n    return"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Now we can run this function and hopefully get the same result.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "reg_func(\"transformed_pipeline.png\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        ".. figure:: transformed_pipeline.png\n   :align: center\n\n   Registration result using pipeline.\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Now let's modify the images in order to test masking. We will place three\nsquares in the corners of both images, but in slightly different locations.\n\nWe will make masks that cover these regions but with an extra border of pixels.\nThis is because the masks need transforming and resampling during optimization,\nand we want to make sure that we are definitely covering the troublesome\nfeatures.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "sz = 15\npd = 10\n\n# modify images\nval = static.max()/2.0\n\nstatic[-sz-pd:-pd, -sz-pd:-pd, :] = val\nstatic[pd:sz+pd, -sz-pd:-pd, :] = val\nstatic[-sz-pd:-pd, pd:sz+pd, :] = val\n\nmoving[pd:sz+pd, pd:sz+pd, :] = val\nmoving[pd:sz+pd, -sz-pd:-pd, :] = val\nmoving[-sz-pd:-pd, pd:sz+pd, :] = val\n\n# create masks\nsquares_st = np.zeros_like(static).astype(np.int32)\nsquares_mv = np.zeros_like(static).astype(np.int32)\n\nsquares_st[-sz-1-pd:-pd, -sz-1-pd:-pd, :] = 1\nsquares_st[pd:sz+1+pd, -sz-1-pd:-pd, :] = 1\nsquares_st[-sz-1-pd:-pd, pd:sz+1+pd, :] = 1\n\nsquares_mv[pd:sz+1+pd, pd:sz+1+pd, :] = 1\nsquares_mv[pd:sz+1+pd, -sz-1-pd:-pd, :] = 1\nsquares_mv[-sz-1-pd:-pd, pd:sz+1+pd, :] = 1\n\n\nregtools.overlay_slices(static, moving, None, 2,\n                        \"Static\", \"Moving\", \"deregistered_squares.png\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        ".. figure:: deregistered_squares.png\n   :align: center\n\n   Same images but misaligned, with white squares in the corners.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "static_mask = np.abs(squares_st - 1)\nmoving_mask = np.abs(squares_mv - 1)\n\nfig, ax = plt.subplots(1, 2)\nax[0].imshow(static_mask[:, :, 1].T, cmap=\"gray\", origin=\"lower\")\nax[0].set_title(\"static image mask\")\nax[1].imshow(moving_mask[:, :, 1].T, cmap=\"gray\", origin=\"lower\")\nax[1].set_title(\"moving image mask\")\nplt.savefig(\"masked_static.png\", bbox_inches='tight')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        ".. figure:: masked_static.png\n   :align: center\n\n   The masks.\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's try to register these new images without a mask.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "reg_func(\"transformed_squares.png\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        ".. figure:: transformed_squares.png\n   :align: center\n\n   Registration fails to align the images because the squares pin the images.\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Now we will attempt to register the images using the masks that we defined.\n\nFirst, use a mask on the static image. Only pixels where the mask is non-zero\nin the static image will contribute to Mutual Information.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "reg_func(\"transformed_squares_mask.png\", static_mask=static_mask)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        ".. figure:: transformed_squares_mask.png\n   :align: center\n\n   Registration result using a static mask.\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We can also attempt the same thing use a moving image mask.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "reg_func(\"transformed_squares_mask_2.png\", moving_mask=moving_mask)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        ".. figure:: transformed_squares_mask_2.png\n   :align: center\n\n   Registration result using a moving mask.\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "And finally, we can use both masks at the same time.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "reg_func(\"transformed_squares_mask_3.png\",\n         static_mask=static_mask, moving_mask=moving_mask)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        ".. figure:: transformed_squares_mask_3.png\n   :align: center\n\n   Registration result using both a static mask and a moving mask.\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "In most use cases, it is likely that only a static mask will be required,\ne.g., to register a series of images to a single static image.\n\nLet's make a series of volumes to demonstrate this idea, and register the\nseries to the first image in the series using a static mask:\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "series = np.stack([static, moving, moving], axis=-1)\n\npipeline = [translation, rigid]\nxformed, _ = register_series(series, 0, pipeline,\n                             series_affine=moving_affine,\n                             static_mask=static_mask)\n\nregtools.overlay_slices(np.squeeze(xformed[..., 0]),\n                        np.squeeze(xformed[..., -2]),\n                        None, 2, \"Static\", \"Moving 1\", \"series_mask_1.png\")\n\nregtools.overlay_slices(np.squeeze(xformed[..., 0]),\n                        np.squeeze(xformed[..., -1]),\n                        None, 2, \"Static\", \"Moving 2\", \"series_mask_2.png\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        ".. figure:: series_mask_1.png\n   :align: center\n.. figure:: series_mask_2.png\n   :align: center\n\n   Registration of series using a static mask.\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "In all of the examples above, different masking choices achieved essentially\nthe same result, but in general the results may differ depending on differences\nbetween the static and moving images.\n\n\n.. [Mattes03] Mattes, D., Haynor, D. R., Vesselle, H., Lewellen, T. K.,\n              Eubank, W. (2003). PET-CT image registration in the chest using\n              free-form deformations. IEEE Transactions on Medical Imaging,\n              22(1), 120-8.\n.. [Avants11] Avants, B. B., Tustison, N., & Song, G. (2011). Advanced\n              Normalization Tools (ANTS), 1-35.\n\n.. include:: ../links_names.inc\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
}