synthmorph_affine_author.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "p-gUO_4ugpPk"
      },
      "outputs": [],
      "source": [
        "# This demo trains an anatomy-aware affine registration network with SynthMorph,\n",
        "# a strategy for learning image registration with wildy variable synthetic data.\n",
        "# Further information is availabe at https://w3id.org/synthmorph.\n",
        "#\n",
        "# If you find our demo useful, please cite:\n",
        "#\n",
        "#     Anatomy-specific acquisition-agnostic affine registration learned from fictitious images\n",
        "#     Hoffmann M, Hoopes A, Fischl B*, Dalca AV* (*equal contribution)\n",
        "#     SPIE Medical Imaging: Image Processing, 12464, p 1246402, 2023\n",
        "#     https://doi.org/10.1117/12.2653251\n",
        "#     https://malte.cz/#papers (PDF)\n",
        "#\n",
        "# We distribute this notebook under the \"2-clause BSD\" license:\n",
        "# https://choosealicense.com/licenses/bsd-2-clause"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "suUJRPSHFDAg",
        "outputId": "5a7c9d75-4be2-495b-a1cf-0f86cdbdaff7"
      },
      "outputs": [],
      "source": [
        "# # Packages from GitHub.\n",
        "# !pip -q install git+https://github.com/adalca/pystrum.git@ba35d4b357f54e5ed577cbd413076a07ef810a21\n",
        "# !pip -q install git+https://github.com/adalca/neurite.git@97ca37940d0c315dc10176e33e57982b3740368e\n",
        "# !pip -q install git+https://github.com/freesurfer/surfa.git@0842291322a05a8fb74f052b881c5679532cb52f\n",
        "# !pip -q install git+https://github.com/voxelmorph/voxelmorph.git@1a2a9529396cdc60cc50c19e0ed4c633ab06589f"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "%load_ext autoreload\n",
        "%autoreload 2"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "p6kExhij7_gs"
      },
      "outputs": [],
      "source": [
        "# Deefault libraries\n",
        "from pathlib import Path\n",
        "\n",
        "# External libraries\n",
        "import cv2\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "import matplotlib.pyplot as plt\n",
        "from skimage import transform\n",
        "import surfa as sf\n",
        "import neurite as ne\n",
        "import voxelmorph as vxm\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Prevent TF model from taking whole GPU memory\n",
        "gpus = tf.config.experimental.list_physical_devices('GPU')\n",
        "for gpu in gpus:\n",
        "        tf.config.experimental.set_virtual_device_configuration(gpu,[tf.config.experimental.VirtualDeviceConfiguration(memory_limit=4096)])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Training"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "q_9orQW_QF-t"
      },
      "outputs": [],
      "source": [
        "# SynthStrip dataset.\n",
        "data = 'synthstrip_data_v1.4_2d'\n",
        "\n",
        "!curl -sO https://surfer.nmr.mgh.harvard.edu/docs/synthstrip/data/{data}.tar\n",
        "!tar xf {data}.tar\n",
        "!ls"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Kb_3yreQRK9d"
      },
      "outputs": [],
      "source": [
        "# Label maps for synthesis.\n",
        "labels, label_maps = vxm.py.utils.load_labels(f'{data}/*/labels.nii.gz')\n",
        "in_shape = label_maps[0].shape\n",
        "\n",
        "\n",
        "# Color map.\n",
        "lut = ne.py.utils.load_fs_lut(f'{data}/seg_labels.txt')\n",
        "cmap = ne.py.utils.fs_lut_to_cmap(lut)\n",
        "\n",
        "\n",
        "# Visualize.\n",
        "num_row = 2\n",
        "per_row = 8\n",
        "for i in range(0, num_row * per_row, per_row):\n",
        "    ne.plot.slices(label_maps[i:i + per_row], cmaps=[cmap])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-98zhiS3kBOR"
      },
      "outputs": [],
      "source": [
        "# Merge labels into larger structures.\n",
        "labels_out = {}\n",
        "for k, v in lut.items():\n",
        "    if 'Cerebellum' in v['name']:\n",
        "        labels_out[k] = 'Cerebellum'\n",
        "    elif v['name'].startswith(('Left', 'Right')):\n",
        "        labels_out[k] = 'Cerebrum'\n",
        "\n",
        "\n",
        "# Show mapping.\n",
        "width = max(len(v['name']) for v in lut.values())\n",
        "for k, v in labels_out.items():\n",
        "    print(lut[k]['name'].rjust(width), '->', v)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "X_p2ZbwxiwUe"
      },
      "outputs": [],
      "source": [
        "# Training-image generation.\n",
        "prop = dict(\n",
        "    in_shape=in_shape,\n",
        "    labels_in=labels,\n",
        "    labels_out=labels_out,\n",
        "    aff_shift=30,\n",
        "    aff_rotate=45,\n",
        "    aff_scale=0.1,\n",
        "    aff_shear=0.1,\n",
        "    blur_max=1,\n",
        "    slice_prob=1,\n",
        "    crop_prob=1,\n",
        ")\n",
        "model_gen_1 = ne.models.labels_to_image_new(**prop, id=1)\n",
        "model_gen_2 = ne.models.labels_to_image_new(**prop, id=2)\n",
        "\n",
        "\n",
        "# Test repeatedly on the same input.\n",
        "num_gen = 8\n",
        "input = np.expand_dims(label_maps[0], axis=(0, -1))\n",
        "images, one_hot = zip(*[model_gen_1.predict(input, verbose=0) for _ in range(num_gen)])\n",
        "\n",
        "\n",
        "# Example images and structures to align.\n",
        "ne.plot.slices(images)\n",
        "for i in range(one_hot[0].shape[-1]):\n",
        "    ne.plot.slices([f[..., i] for f in one_hot])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3br8h7eN-3Fs"
      },
      "outputs": [],
      "source": [
        "# Registration model.\n",
        "model_aff = vxm.networks.VxmAffineFeatureDetector(in_shape)\n",
        "\n",
        "\n",
        "# Combined model: synthesis and registration.\n",
        "ima_1, map_1 = model_gen_1.outputs\n",
        "ima_2, map_2 = model_gen_2.outputs\n",
        "\n",
        "trans = model_aff((ima_1, ima_2))\n",
        "moved = vxm.layers.SpatialTransformer(fill_value=0)((map_1, trans))\n",
        "\n",
        "inputs = (*model_gen_1.inputs, *model_gen_2.inputs)\n",
        "model = tf.keras.Model(inputs, outputs=(moved, trans))\n",
        "\n",
        "\n",
        "# Contrast invariance: MSE loss on probability maps.\n",
        "model.add_loss(vxm.losses.MSE().loss(map_2, moved))\n",
        "model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=1e-4))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ez7RywMN-3cs"
      },
      "outputs": [],
      "source": [
        "# Training. Re-running the cell will continue training.\n",
        "hist = model.fit(\n",
        "    x=vxm.generators.synthmorph(label_maps),\n",
        "    epochs=3,\n",
        "    steps_per_epoch=100,\n",
        ")\n",
        "\n",
        "\n",
        "# Visualize loss.\n",
        "plt.plot(hist.epoch, hist.history['loss'], '.-')\n",
        "plt.xlabel('Epoch')\n",
        "plt.ylabel('Loss');"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "h4IDUfnz-3sU",
        "outputId": "33297bc7-bdd0-43ab-f569-a6bcaf453516"
      },
      "outputs": [],
      "source": [
        "# Skip training, download model weights.\n",
        "!gdown -O weights.h5 1DWiVxCvQmYDSBS1RcVbUeTxX0XArQsGv\n",
        "# model.load_weights('weights.h5')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dxXrMKNl-4Cj"
      },
      "outputs": [],
      "source": [
        "# Potential test subjects. Training used label maps from T1w ASL/FSM/IXI scans.\n",
        "%ls {data}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Testing"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "spmfHHiqWEoM"
      },
      "outputs": [],
      "source": [
        "# Conform input images.\n",
        "def conform(f):\n",
        "    # Conform to shape, isotropic 1-mm resolution. The function interpolates in\n",
        "    # 3D, so we stack 3 copies of the 2D image to avoid all-zero outputs when\n",
        "    # the through-plane coordinate is not exactly zero.\n",
        "    out = sf.load_volume(f)\n",
        "    pad = np.concatenate([out.data] * 3, axis=-1)\n",
        "    out = out.new(pad)\n",
        "    out = out.conform(\n",
        "        voxsize=(1, 1, out.geom.voxsize[2]),\n",
        "        shape=(*in_shape, out.shape[-1]),\n",
        "        dtype=np.float32,\n",
        "        method='nearest',\n",
        "    )[..., 1]\n",
        "\n",
        "    # Normalize, add batch and feature dimension.\n",
        "    out = ne.utils.minmax_norm(out)\n",
        "    return out[None, ..., None]\n",
        "\n",
        "\n",
        "# Skull-strip for difference images.\n",
        "def load_and_strip(subj):\n",
        "    imag = conform(f=f'{data}/{subj}/image.nii.gz')\n",
        "    mask = conform(f=f'{data}/{subj}/mask.nii.gz')\n",
        "    return imag, tf.multiply(imag, mask)\n",
        "\n",
        "\n",
        "# # Visualize output.\n",
        "# slices = load_and_strip(subj='fsm_qt1_79bf')\n",
        "# ne.plot.slices(slices, titles=('Full image', 'Skull-stripped'), width=len(slices) * 3);"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4m9bHgUyPwCg"
      },
      "outputs": [],
      "source": [
        "def generate_trapezoid(image_size):\n",
        "    # Create a blank image\n",
        "    image = np.zeros(image_size)\n",
        "\n",
        "    # Define the coordinates of the rectangle\n",
        "    midpoint = image_size[0] // 2\n",
        "    x_topleft = midpoint - (midpoint // 3)\n",
        "    x_topright = midpoint + (midpoint // 3)\n",
        "    x_bottomleft =  midpoint - (midpoint // 2) - (midpoint // 8)\n",
        "    x_bottomright = midpoint + (midpoint // 2) + (midpoint // 8)\n",
        "    y_top = midpoint - (midpoint // 3)\n",
        "    y_bottom = midpoint + (midpoint // 3)\n",
        "    rectangle_coords = np.array([[x_topleft, y_top], [x_topright, y_top], [x_bottomright, y_bottom], [x_bottomleft, y_bottom]])\n",
        "\n",
        "    # Fill the rectangle region with white color (1.0)\n",
        "    image = cv2.fillPoly(image, [rectangle_coords], 1.0)\n",
        "\n",
        "    return image\n",
        "\n",
        "def shift(image, translation):\n",
        "    # Generate an affine transformation matrix\n",
        "    affine_matrix = np.array([[1, 0, translation[0]], [0, 1, translation[1]], [0, 0, 1]])\n",
        "\n",
        "    # Apply the affine transformation\n",
        "    transformed_image = transform.warp(image, inverse_map=affine_matrix)\n",
        "\n",
        "    return transformed_image"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9537CbaeOppZ"
      },
      "outputs": [],
      "source": [
        "keras_weights_path = Path(\".\") / \"weights\" / \"keras\"\n",
        "affine_weights = keras_weights_path / \"affine_author_2d_256.h5\"\n",
        "in_shape = (256,) * 2\n",
        "model_aff = vxm.networks.VxmAffineFeatureDetector(in_shape)\n",
        "model_aff.load_weights(affine_weights)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "model_aff.summary()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nXRU3E0eKcTz"
      },
      "outputs": [],
      "source": [
        "minmax_norm = lambda x: ne.utils.minmax_norm(x)\n",
        "translation = (32, 32)\n",
        "\n",
        "moving = generate_trapezoid(in_shape)\n",
        "fixed = shift(moving, translation)\n",
        "moving, fixed = minmax_norm(np.expand_dims(moving, (-1, 0,))), minmax_norm(np.expand_dims(fixed, (-1, 0)))\n",
        "trans = model_aff((moving, fixed))\n",
        "print(f\"trans: {[t.shape for t in trans]}\")\n",
        "moved = vxm.layers.SpatialTransformer(fill_value=0)((moving, trans))\n",
        "\n",
        "# Display the images\n",
        "plt.subplot(1, 3, 1)\n",
        "plt.imshow(moving.numpy().squeeze(), cmap='gray')\n",
        "plt.title('Moving')\n",
        "\n",
        "plt.subplot(1, 3, 2)\n",
        "plt.imshow(fixed.numpy().squeeze(), cmap='gray')\n",
        "plt.title('Fixed')\n",
        "\n",
        "plt.subplot(1, 3, 3)\n",
        "plt.imshow(minmax_norm(moved.numpy().squeeze()), cmap='gray')\n",
        "plt.title('Moved')\n",
        "plt.tight_layout()\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ik3584sZ7gO7"
      },
      "outputs": [],
      "source": [
        "\n",
        "# Test registration.\n",
        "def register(moving, fixed):\n",
        "    # Load images.\n",
        "    im_1, br_1 = load_and_strip(moving)\n",
        "    im_2, br_2 = load_and_strip(fixed)\n",
        "\n",
        "    # Register.\n",
        "    trans = model_aff.predict((im_1, im_2), verbose=0)\n",
        "    out_im = vxm.layers.SpatialTransformer(fill_value=0)((im_1, trans))\n",
        "    out_br = vxm.layers.SpatialTransformer(fill_value=0)((br_1, trans))\n",
        "\n",
        "    # Re-normalize for visualization.\n",
        "    slices = (im_1, im_2, br_1, br_2, out_im, out_br)\n",
        "    slices = (np.clip(f, *np.percentile(f, q=(0.1, 99.9))) for f in slices)\n",
        "    im_1, im_2, br_1, br_2, out_im, out_br = slices\n",
        "\n",
        "    # Plot.\n",
        "    slices = (im_1, im_2, out_im, br_1 - br_2, out_br - br_2)\n",
        "    titles = ('Moving', 'Fixed', 'Moved', 'Difference before', 'Difference after')\n",
        "    ne.plot.slices(slices, titles=titles, width=len(slices) * 3)\n",
        "\n",
        "\n",
        "# Defaced quantitative-T1 map to low-resolution DWI.\n",
        "register(moving='fsm_qt1_79bf', fixed='ixi_dwi_401')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QoGdsTbF0mLA"
      },
      "outputs": [],
      "source": [
        "# Partial-FOV MRA to different-partial-FOV MRA.\n",
        "register(moving='ixi_mra_012', fixed='ixi_mra_016')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ai6PWjNa05Db"
      },
      "outputs": [],
      "source": [
        "# Clinical thick-slice T2-FLAIR 2D-FSE to defaced PDw 3D-FLASH.\n",
        "register(moving='qin_flair_07', fixed='fsm_pd_50ww')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IXjh7anY5SO4"
      },
      "outputs": [],
      "source": [
        "# Low-resolution partial-FOV 2D-EPI to defaced quantitative-T1 map.\n",
        "register(moving='asl_epi_134', fixed='fsm_qt1_87qb')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sQTXB6o1BI_M"
      },
      "outputs": [],
      "source": [
        "# Clinical T2w thick-slice 2D-FSE with glitch to T1w infant MPRAGE.\n",
        "register(moving='qin_t2_05', fixed='infant_t1_17')"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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.8.18"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}