run_OCRT.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Choose a sample to run. `sample_id` can be one of:\n",
    "- mouse_vas_deferens1\n",
    "- mouse_vas_deferens2\n",
    "- mouse_femoral_artery\n",
    "- mouse_bladder\n",
    "- mouse_trachea\n",
    "- human_cornea\n",
    "- insect_leg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_id = 'human_cornea'  # a string that denotes which data to run\n",
    "data_directory = 'data/'  # where is the OCRT input data located?\n",
    "save_directory = 'saved_models_and_variables/'  # where to save the tf graph after optimization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Registration of multiangle B-scans and synthesis of refractive index map\n",
    "The registration metric is optimized with respect to the deformation model, which in this case is ray propagation through a spatially inhomogeneous refractive index map."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function, division\n",
    "from OCRT import OCRT2D\n",
    "import tensorflow.compat.v1 as tf\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from time import time\n",
    "%matplotlib inline\n",
    "\n",
    "a = OCRT2D(sample_id=sample_id,save_directory=save_directory)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# adjust some parameters, depending on the sample:\n",
    "if sample_id in ['mouse_bladder', 'mouse_femoral_artery', 'mouse_trachea', 'mouse_vas_deferens1']:\n",
    "    num_iter = 200\n",
    "elif sample_id in ['mouse_vas_deferens2']:\n",
    "    num_iter = 300\n",
    "elif sample_id in ['insect_leg']:\n",
    "    num_iter = 500\n",
    "elif sample_id in ['human_cornea']:\n",
    "    num_iter = 500\n",
    "    # this sample is more difficult to register, so use a multiresolution approach:\n",
    "    a.use_multires = True\n",
    "    a.size_factor_ = 1  # the final size_factor_ is 8\n",
    "    a.switch_iter = 250\n",
    "else:\n",
    "    raise Exception('invalid sample_id: ' + sample_id)\n",
    "    \n",
    "# two tube sizes were used (same inner diameter, but different outer diameter):\n",
    "if sample_id in ['mouse_vas_deferens1', 'mouse_vas_deferens2', 'mouse_trachea']:\n",
    "    a.tube_diameter = 1.066  # in mm\n",
    "else:\n",
    "    a.tube_diameter = 1.108516"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a.load_data_and_resolve_constants(data_directory=data_directory)\n",
    "a.build_graph()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# run through optimization loop:\n",
    "losses = list()\n",
    "feed_dict = a.get_feed_dict()\n",
    "for i in range(num_iter + 1):\n",
    "    \n",
    "    # if using multires, change the pixel resolution of the reconstruction\n",
    "    if i == a.switch_iter and a.use_multires: \n",
    "        feed_dict[a.size_factor] = a.final_size_factor\n",
    "    \n",
    "    start = time()\n",
    "    loss_i, _ = a.sess.run([a.loss_terms, a.train_op], feed_dict=feed_dict)\n",
    "    losses.append(loss_i)\n",
    "    print(i, loss_i, time()-start)\n",
    "    # loss_i is a list of all the contributors to the scalar loss;\n",
    "    # (see a.loss_terms or a.loss_term_names for names of the regularization terms)\n",
    "    \n",
    "    # monitor results periodically:\n",
    "    if i % 10 == 0:\n",
    "        recon_i = a.sess.run(a.recon, feed_dict=feed_dict)\n",
    "        recon_i = recon_i.sum(2)  # only once slice along y contains nonzero values, because we are optimizing 2D\n",
    "        plt.figure(figsize=(10, 10))\n",
    "        plt.imshow(recon_i, cmap='gray_r')\n",
    "        plt.title('OCRT reconstruction')\n",
    "        plt.show()\n",
    "        \n",
    "        RI = a.sess.run(a.RI, feed_dict={a.xz_delta: np.zeros((60, 2))})  # remove xy_delta shifts\n",
    "        \n",
    "        plt.imshow(RI)\n",
    "        plt.title('refractive index map')\n",
    "        plt.colorbar()\n",
    "        plt.show()\n",
    "        \n",
    "        plt.plot(losses)\n",
    "        plt.legend(a.loss_term_names.eval())\n",
    "        plt.title('loss terms')\n",
    "        plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a.save_graph()  # this graph must be saved in order to run filter optimization below"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Filter optimization after registration\n",
    "Freeze the registration/refractive index parameters and only optimize the 2D backprojection filter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# remove previous tf graph and variables:\n",
    "a.sess.close()\n",
    "tf.reset_default_graph()\n",
    "del a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# instantiate new object, this time for filter optimization\n",
    "a = OCRT2D(sample_id=sample_id,save_directory=save_directory)\n",
    "num_iter = 100\n",
    "a.infer_backprojection_filter = True\n",
    "a.use_spatial_shifts = False\n",
    "\n",
    "# as above, set the tube diameter depending on the sample:\n",
    "if sample_id in ['mouse_vas_deferens1', 'mouse_vas_deferens2', 'mouse_trachea']:\n",
    "    a.tube_diameter = 1.066\n",
    "else:\n",
    "    a.tube_diameter = 1.108516\n",
    "\n",
    "a.load_data_and_resolve_constants(data_directory=data_directory)\n",
    "a.build_graph()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "losses = list()\n",
    "feed_dict = a.get_feed_dict()\n",
    "for i in range(num_iter + 1):\n",
    "    start = time()\n",
    "    loss_i, _ = a.sess.run([a.loss_terms, a.train_op], feed_dict=feed_dict)\n",
    "    losses.append(loss_i)\n",
    "    print(i, loss_i, time()-start)\n",
    "    \n",
    "    # monitor results periodically:\n",
    "    if i % 20 == 0:\n",
    "        recon_i = a.sess.run(a.recon, feed_dict=feed_dict)\n",
    "        recon_i = recon_i.sum(2)\n",
    "        plt.figure(figsize=(10, 10))\n",
    "        plt.imshow(recon_i, cmap='gray_r')\n",
    "        plt.title('OCRT reconstruction')\n",
    "        plt.show()\n",
    "        \n",
    "        plt.plot(losses)\n",
    "        plt.legend(a.loss_term_names.eval())\n",
    "        plt.title('loss terms')\n",
    "        plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# a.save_graph()  # this graph doesn't need to be saved"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tf1.8",
   "language": "python",
   "name": "tf1.8"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}