run_mesoscopic_photogrammetry.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Overview\n",
    "There are 3 rounds of optimization at different resolutions. The first two rounds are common to all methods. There are two separate sections for round 3, depending on whether to regularize with CNN reparameterization or total variation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import os\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '0'  # make only one GPU visible if multiple are present\n",
    "os.environ['TF_GPU_HOST_MEM_LIMIT_IN_MB'] = '32000'  # for tensorflow LMS\n",
    "import matplotlib.pyplot as plt\n",
    "import cv2\n",
    "from time import time\n",
    "import tensorflow as tf\n",
    "from mesoSfM import mesoSfM, stack_loader_phone, monitor_progress, xcorr_initial_guess\n",
    "from tqdm.notebook import tqdm\n",
    "from datetime import date\n",
    "\n",
    "tf.config.experimental.set_lms_enabled(True)\n",
    "tf.config.experimental.set_lms_defrag_enabled(True)\n",
    "\n",
    "def set_seed(seed=0):\n",
    "    tf.random.set_seed(seed)  # this doesn't deal with nondeterministic GPU operations though\n",
    "    np.random.seed(seed)\n",
    "set_seed(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Specify dataset and hyperparameters\n",
    "Specify the dataset and distortion method in the next cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = 'cut_cards'  # one of the following: cut_cards, painting, helicopter_seeds, PCB, tuning_set\n",
    "distort_method = 'piecewise_linear'  # 'radial' or 'piecewise_linear'\n",
    "num_radial_terms = 2  # only used for 'radial' distort model; the max power is twice this number;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "directory = './data/' + dataset + '/'\n",
    "im_stack = stack_loader_phone(directory)  # load images into memory\n",
    "num_channels = im_stack.shape[-1]\n",
    "inds_keep = np.arange(len(im_stack))  # which images to use, specified by indices\n",
    "recon_shape = np.array((5500, 6500))  # shape of reconstruction in pixels\n",
    "dither_coords = True\n",
    "\n",
    "# if difficult to get initial registration, use sequential cross-correlation-based estimate:\n",
    "if dataset in ['helicopter_seeds', 'tuning_set', 'PCB']:\n",
    "    x_pos, y_pos = xcorr_initial_guess(im_stack)\n",
    "elif dataset in ['cut_cards']:\n",
    "    x_pos, y_pos = xcorr_initial_guess(im_stack, crop_frac=.2)\n",
    "else:\n",
    "    # otherwise, just initialize initial positions with 0s;\n",
    "    x_pos = np.zeros(len(im_stack))\n",
    "    y_pos = np.zeros(len(im_stack))\n",
    "    \n",
    "# global xy offset in reconstruction in pixels;\n",
    "ul_offset = np.array([1100, 1100])\n",
    "if 'PCB' in directory:\n",
    "    ul_offset = np.array([1400, 1400])\n",
    "\n",
    "# pre-calibrated magnifications using control points:\n",
    "if 'cut_cards' in directory:\n",
    "    magnification = 0.069315160567588\n",
    "elif 'tuning_set' in directory:\n",
    "    magnification = 0.0835\n",
    "elif 'painting' in directory:\n",
    "    magnification = 0.058971786873920\n",
    "elif 'PCB' in directory:\n",
    "    magnification = 0.076508519683842\n",
    "elif 'helicopter_seeds' in directory:\n",
    "    magnification = 0.076790599352193\n",
    "else:\n",
    "    raise Exception('magnification not defined')\n",
    "\n",
    "# I already downsampled the images by 2x, so indicate that here:\n",
    "pre_downsample_factor = 2\n",
    "\n",
    "recon_shape = recon_shape // pre_downsample_factor\n",
    "ul_offset = ul_offset // pre_downsample_factor"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Round 1 of optimization\n",
    "First, optimize at a low resolution (heavy downsampling), updating only the xy shifts and keeping everything else fixed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "set_seed(1)\n",
    "\n",
    "# create mesoSfM object with particular settings:\n",
    "a = mesoSfM(stack=im_stack[inds_keep], ul_coords=np.stack([y_pos[inds_keep], x_pos[inds_keep]]).T,\n",
    "            recon_shape=recon_shape, ul_offset=ul_offset, scale=.05*pre_downsample_factor)\n",
    "\n",
    "# create tf variables:\n",
    "a.create_variables(deformation_model='camera_parameters_perspective_to_orthographic',\n",
    "    learning_rates={'camera_focal_length': -1e-3, 'camera_height': -1e-3, 'ground_surface_normal': -1e-3,\n",
    "                   'camera_in_plane_angle': -1e-3, 'rc': 10, 'gain': -1e-3, 'bias':-1e-3, 'ego_height': -1e-3},\n",
    "                   variable_initial_values=None, remove_global_transform=True, antialiasing_filter=True)\n",
    "\n",
    "# create dataset:\n",
    "stack_downsamp, rc_downsamp = a.generate_dataset()\n",
    "\n",
    "# optimization loop:\n",
    "losses = list()\n",
    "for ii in tqdm(range(401)):\n",
    "    start = time()\n",
    "    loss_i, recon, normalize, error_map = a.gradient_update(stack_downsamp, rc_downsamp)\n",
    "    if type(loss_i) is list:\n",
    "        losses.append([loss.numpy() for loss in loss_i])\n",
    "    else:\n",
    "        losses.append(loss_i.numpy())\n",
    "    \n",
    "    if ii % 10 == 0:\n",
    "        print(ii, losses[-1], time() - start)\n",
    "    if ii % 100 == 0:\n",
    "        monitor_progress(recon, error_map, losses)\n",
    "        \n",
    "variable_initial_values = a.get_all_variables()  # store the optimized variables to initialize the next round;"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Round 2 of optimization\n",
    "Next, optimize at a higher resolution, updating all parameters, including distortion, except for height map."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "set_seed(2)\n",
    "\n",
    "learning_rates = {'camera_focal_length': -1e-3, 'camera_height': 1e-3, 'ground_surface_normal': 1e-3,\n",
    "                   'camera_in_plane_angle': 1e-3, 'rc': .1, 'gain': -1e-3, 'bias': -1e-3, 'ego_height': -1e-3}\n",
    "                   \n",
    "# initialization and learning rates for distortion models:\n",
    "if distort_method == 'piecewise_linear':\n",
    "    variable_initial_values['camera_distortion_center'] = np.zeros([1, 2])+1e-7\n",
    "    variable_initial_values['radial_camera_distortion_piecewise_linear'] = np.zeros(50)\n",
    "    learning_rates['radial_camera_distortion_piecewise_linear'] = 1e-3\n",
    "    learning_rates['camera_distortion_center'] = 1e-4\n",
    "elif distort_method == 'radial':\n",
    "    variable_initial_values['camera_distortion_center'] = np.zeros([1, 2])+1e-7\n",
    "    variable_initial_values['radial_camera_distortion'] = np.zeros([1, num_radial_terms])\n",
    "    learning_rates['camera_distortion_center'] = 1e-4\n",
    "    learning_rates['radial_camera_distortion'] = 1e-2\n",
    "else:\n",
    "    raise Exception('invalid distort_method')\n",
    "\n",
    "# create mesoSfM object with particular settings:\n",
    "a = mesoSfM(stack=im_stack[inds_keep], ul_coords=np.stack([y_pos[inds_keep], x_pos[inds_keep]]).T,\n",
    "            recon_shape=recon_shape, ul_offset=ul_offset, scale=.15*pre_downsample_factor)\n",
    "\n",
    "# create tf Variables:\n",
    "a.create_variables(deformation_model='camera_parameters_perspective_to_orthographic',\n",
    "                   learning_rates=learning_rates,\n",
    "                   variable_initial_values=variable_initial_values,\n",
    "                   remove_global_transform=True, antialiasing_filter=True)\n",
    "\n",
    "# create dataset:\n",
    "stack_downsamp, rc_downsamp = a.generate_dataset()  # when not batching, this is just a one-batch dataset;\n",
    "\n",
    "# optimization loop:\n",
    "losses = list()\n",
    "for ii in tqdm(range(401)):\n",
    "    start = time()\n",
    "    loss_i, recon, normalize, error_map = a.gradient_update(stack_downsamp, rc_downsamp)\n",
    "    if type(loss_i) is list:\n",
    "        losses.append([loss.numpy() for loss in loss_i])\n",
    "    else:\n",
    "        losses.append(loss_i.numpy())\n",
    "    if ii % 10 == 0:\n",
    "        print(ii, losses[-1], time() - start)\n",
    "        \n",
    "    if ii % 100 == 0:\n",
    "        monitor_progress(recon, error_map, losses)\n",
    "        if len(losses) > 201:\n",
    "            plt.plot(losses[200:])\n",
    "            plt.show()\n",
    "            \n",
    "variable_initial_values = a.get_all_variables()  # store the optimized variables to initialize the next round;\n",
    "variable_indices = {var.name: i for i, var in enumerate(a.train_var_list)}  # to identify the optimizers\n",
    "# so that we can anneal the learning rate (for CNN);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Two options from this point\n",
    "Now that we've gotten a decent optimized result using only camera parameters, we're ready to go to the final resolution AND optimize the height map. We have two options for round 3 for regularizing the height map: 1) reparameterization with a convolutional neural network (CNN), and 2) total variation (TV). Run only the cells corresponding to the desired regularization -- these cells are separated by different headers:\n",
    "- Round 3 of optimization with total variation\n",
    "- Round 3 of optimization with CNN reparameterization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Round 3 of optimization with CNN reparameterization\n",
    "To test out different CNN architectures, modify `CNN_architecture` below, which is a list of filter numbers (see supplementary document of paper). The main one used for the paper was `[16, 16, 16, 32, 32]`, but other ones tested were: `[16, 16, 32, 32]`, `[16, 16, 16, 16]`, and `[16, 16, 16, 16, 16]`. This round will take up the bulk of the total processing time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "CNN_architecture = [16, 16, 16, 32, 32]\n",
    "reg_coefs = {'height_map': .1}\n",
    "num_iter = 10001\n",
    "\n",
    "tf.config.optimizer.set_experimental_options({'layout_optimizer': False})  # might save some memory\n",
    "\n",
    "# specify learning rates:\n",
    "learning_rates = {'camera_focal_length': -1e-3, 'camera_height': 1e-3, 'ground_surface_normal': 1e-3,\n",
    "                   'camera_in_plane_angle': 1e-3, 'rc': .1, 'gain': -1e-3, 'bias': -1e-3, 'ego_height': 1e-3}\n",
    "if distort_method == 'piecewise_linear':\n",
    "    learning_rates['radial_camera_distortion_piecewise_linear'] = 1e-3\n",
    "    learning_rates['camera_distortion_center'] = 1e-4\n",
    "elif distort_method == 'radial':\n",
    "    learning_rates['camera_distortion_center'] = 1e-4\n",
    "    learning_rates['radial_camera_distortion'] = 1e-2\n",
    "else:\n",
    "    raise Exception('invalid distort_method')\n",
    "    \n",
    "# criterion for optimization divergence, to roll back to a previous checkpoint:\n",
    "height_map_reg_threshold = 8\n",
    "    \n",
    "if distort_method == 'radial':\n",
    "    height_map_reg_threshold *= 2  # with a poorer distortion model, the height maps are less consistent;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "seed = 5\n",
    "while True:\n",
    "    # sometimes, the optimization unpredictably encounters an NaN, or it doesn't start converging; \n",
    "    # if this happens start over;\n",
    "    # the try-except in the while loop basically says keep trying to complete the full optimization loop;\n",
    "    \n",
    "    try:\n",
    "        set_seed(seed)\n",
    "\n",
    "        # create mesoSfM object with particular settings:\n",
    "        a = mesoSfM(stack=im_stack[inds_keep], ul_coords=np.stack([y_pos[inds_keep], x_pos[inds_keep]]).T,\n",
    "                    recon_shape=recon_shape, ul_offset=ul_offset, scale=.5*pre_downsample_factor,\n",
    "                    batch_size=6, momentum=.5, batch_across_images=True)\n",
    "\n",
    "        # these are for getting absolute height values:\n",
    "        a.use_absolute_scale_calibration = True\n",
    "        a.effective_focal_length_mm = 4.3\n",
    "        a.magnification_j = magnification\n",
    "\n",
    "        # CNN settings:\n",
    "        a.filters_list = CNN_architecture\n",
    "        a.skip_list = [0] * len(CNN_architecture)  # no skip connections\n",
    "        a.save_iter = 50  # checkpoint parameters every this many iterations\n",
    "        a.recompute_CNN = True  # gradient checkpointing to reduce memory usage\n",
    "        \n",
    "        # create tf Variables:\n",
    "        a.create_variables(deformation_model='camera_parameters_perspective_to_orthographic_unet',\n",
    "                           learning_rates=learning_rates, variable_initial_values=variable_initial_values,\n",
    "                           remove_global_transform=True, antialiasing_filter=False)\n",
    "        \n",
    "        # create dataset:\n",
    "        dataset = a.generate_dataset()\n",
    "\n",
    "        # optimization loop:\n",
    "        ii = 0\n",
    "        losses = list()\n",
    "        for stack_downsamp, rc_downsamp in dataset:\n",
    "            start = time()\n",
    "\n",
    "            loss_i = a.gradient_update(stack_downsamp, rc_downsamp, dither_coords=dither_coords,\n",
    "                                       reg_coefs=reg_coefs, return_loss_only=True)\n",
    "            if type(loss_i) is list:\n",
    "                losses.append([loss.numpy() for loss in loss_i])\n",
    "            else:\n",
    "                losses.append(loss_i.numpy())\n",
    "\n",
    "            if ii % 250 == 0:\n",
    "                print(ii, losses[-1], time() - start)\n",
    "            \n",
    "            # conditions for rolling back to an earlier checkpoint:\n",
    "            loss_num = 1\n",
    "            last_few_losses = [loss[loss_num] for loss in losses[-10:-1]]\n",
    "            roll_back_condition = losses[-1][loss_num] > height_map_reg_threshold\n",
    "\n",
    "            if len(losses) <= a.save_iter:\n",
    "                # don't do anything early on, let the optimizer explore a bit, while  checkpoints are accumulating;\n",
    "                pass  \n",
    "\n",
    "            elif (roll_back_condition and (2*a.save_iter >= len(losses) > a.save_iter)) or any(np.isnan(loss_i)):\n",
    "                # now that the optimizer has explored a bit, if it's too far off track early on,\n",
    "                # it may never find its way -- just restart; this exception will be picked up\n",
    "                # by the try-except construct;\n",
    "                seed += 1  # alter the seed to avoid running into the same issue;\n",
    "                raise ValueError('Optimization diverged early on, restarting ...')\n",
    "\n",
    "            elif ii == 2000 and losses[-1][loss_num] < 1e-4:\n",
    "                # if by iteration 2000, the height map reg hasn't attained a certain value, then\n",
    "                # restart;\n",
    "                raise ValueError('Optimization never converged, restarting ...')\n",
    "\n",
    "            elif roll_back_condition and len(losses) > a.save_iter*2:  \n",
    "                # after 2*save_iter, there are two checkpoints to use;\n",
    "                print('Last few losses: ' + str(last_few_losses))\n",
    "                print('Current loss: ' + str(losses[-1][loss_num]))\n",
    "                a.restore_all_variables()\n",
    "                # anneal learning rate:\n",
    "                optim = a.optimizer_list[variable_indices['ego_height:0']]\n",
    "                optim.lr.assign(optim.lr * .9)\n",
    "                # remove most recent loss values so that you don't influence the next test for divergence:\n",
    "                losses = losses[:-2]\n",
    "\n",
    "            if len(losses) % a.save_iter == 0:\n",
    "                # periodically checkpoint all variables for rolling back; see above;\n",
    "                a.checkpoint_all_variables(path='./tf_ckpts')\n",
    "\n",
    "            # break out of for-loop when you reach num_iter iterations:\n",
    "            if ii == num_iter:\n",
    "                break\n",
    "            else:\n",
    "                ii += 1\n",
    "\n",
    "        break\n",
    "    except ValueError:\n",
    "        print('Optimization diverged early on; restarting the optimization ...')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# when all is done, plot results:\n",
    "loss_i, recon, normalize, error_map, tracked = a.gradient_update(stack_downsamp, rc_downsamp, update_gradient=False,\n",
    "                                                                 dither_coords=False, reg_coefs=reg_coefs,\n",
    "                                                                 return_tracked_tensors=True)\n",
    "monitor_progress(recon, error_map, losses, tracked)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Round 3 of optimization with total variation\n",
    "To obtain results with total variation (TV) regularization, after round 2 above, skip down to this cell without running the cells under **Round 3 of optimization with CNN reparameterization**.\n",
    "\n",
    "This round will be done in two steps: \n",
    "1. Optimize all parameters with downsampling. The difference betwen this and round 2 is that the height map is also optimized.\n",
    "2. Optimize all parameters at the final resolution. This step also requires batching to satisfy memory constraints."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### First, optimize all parameters with downsampling\n",
    "Regardless of the final TV coefficient, this step is always run. No need to modify anything here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.config.optimizer.set_experimental_options({'layout_optimizer': False})  # might save some memory\n",
    "set_seed(2)\n",
    "\n",
    "# specify learning rates:\n",
    "learning_rates = {'camera_focal_length': -1e-3, 'camera_height': 1e-3, 'ground_surface_normal': 1e-3,\n",
    "                   'camera_in_plane_angle': 1e-3, 'rc': .1, 'gain': -1e-3, 'bias': -1e-3, 'ego_height': 1e-3}\n",
    "reg_coefs = {'TV': .001, 'height_map': .1}  # specify regularization coefficients;\n",
    "                   \n",
    "if distort_method == 'piecewise_linear':\n",
    "    learning_rates['radial_camera_distortion_piecewise_linear'] = 1e-3\n",
    "    learning_rates['camera_distortion_center'] = 1e-4\n",
    "elif distort_method == 'radial':\n",
    "    learning_rates['camera_distortion_center'] = 1e-4\n",
    "    learning_rates['radial_camera_distortion'] = 1e-2\n",
    "else:\n",
    "    raise Exception('invalid distort_method')\n",
    "\n",
    "# create mesoSfM object with particular settings:\n",
    "a = mesoSfM(stack=im_stack[inds_keep], ul_coords=np.stack([y_pos[inds_keep], x_pos[inds_keep]]).T,\n",
    "            recon_shape=recon_shape, ul_offset=ul_offset, scale=.15*pre_downsample_factor)\n",
    "\n",
    "# these are for getting absolute height values:\n",
    "a.use_absolute_scale_calibration = True\n",
    "a.effective_focal_length_mm = 4.3\n",
    "a.magnification_j = magnification\n",
    "\n",
    "# create tf Variables:\n",
    "a.create_variables(deformation_model='camera_parameters_perspective_to_orthographic',\n",
    "                   learning_rates=learning_rates, variable_initial_values=variable_initial_values,\n",
    "                   remove_global_transform=True, antialiasing_filter=True)\n",
    "\n",
    "# create dataset:\n",
    "stack_downsamp, rc_downsamp = a.generate_dataset()\n",
    "\n",
    "# optimization loop:\n",
    "losses = list()\n",
    "for ii in tqdm(range(401)):\n",
    "    start = time()\n",
    "    loss_i, recon, normalize, error_map, tracked = a.gradient_update(stack_downsamp, rc_downsamp,\n",
    "                                                            reg_coefs=reg_coefs, return_tracked_tensors=True)\n",
    "    if type(loss_i) is list:\n",
    "        losses.append([loss.numpy() for loss in loss_i])\n",
    "    else:\n",
    "        losses.append(loss_i.numpy())\n",
    "    if ii % 10 == 0:\n",
    "        print(ii, losses[-1], time() - start)\n",
    "        \n",
    "    if ii % 100 == 0:\n",
    "        monitor_progress(recon, error_map, losses, tracked)\n",
    "        if len(losses) > 201:\n",
    "            plt.plot(losses[200:])\n",
    "            plt.show()\n",
    "        \n",
    "variable_initial_values = a.get_all_variables()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Second, optimize all parameters at the final resolution\n",
    "Set the desired TV coefficient in the `reg_coefs` dictionary below. The ones used in the paper were 0.003, 0.01, 0.03, and 0.1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "reg_coefs = {'TV': .1, 'height_map': .1}  # set TV reg coef\n",
    "num_iter = 2001  # number of iterations in optimization loop\n",
    "\n",
    "# initialize with previous stitched result:\n",
    "recon_previous = np.copy(recon.numpy())\n",
    "normalize_previous = np.copy(normalize.numpy())\n",
    "\n",
    "learning_rates = {'camera_focal_length': -1e-3, 'camera_height': 1e-3, 'ground_surface_normal': 1e-3,\n",
    "                   'camera_in_plane_angle': 1e-3, 'rc': .1, 'gain': -1e-3, 'bias': -1e-3, 'ego_height': 1e-3}\n",
    "\n",
    "if distort_method == 'piecewise_linear':\n",
    "    learning_rates['radial_camera_distortion_piecewise_linear'] = 1e-3\n",
    "    learning_rates['camera_distortion_center'] = 1e-4\n",
    "elif distort_method == 'radial':\n",
    "    learning_rates['camera_distortion_center'] = 1e-4\n",
    "    learning_rates['radial_camera_distortion'] = 1e-2\n",
    "else:\n",
    "    raise Exception('invalid distort_method')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "while True:\n",
    "    # sometimes, the optimization unpredictably encounters an NaN; if this happens start over;\n",
    "    # the try-except in the while loop basically says keep trying to complete the full optimization loop;\n",
    "    \n",
    "    try:\n",
    "        set_seed(3)\n",
    "\n",
    "        # create mesoSfM object with particular settings:\n",
    "        a = mesoSfM(stack=im_stack[inds_keep], ul_coords=np.stack([y_pos[inds_keep], x_pos[inds_keep]]).T,\n",
    "                    recon_shape=recon_shape, ul_offset=ul_offset, scale=.5*pre_downsample_factor,\n",
    "                    batch_size=6, momentum=.5, batch_across_images=True)\n",
    "        \n",
    "        # these are for getting absolute height values:\n",
    "        a.use_absolute_scale_calibration = True\n",
    "        a.effective_focal_length_mm = 4.3\n",
    "        a.magnification_j = magnification\n",
    "\n",
    "        # create tf Variables:\n",
    "        a.create_variables(deformation_model='camera_parameters_perspective_to_orthographic',\n",
    "                           learning_rates=learning_rates, variable_initial_values=variable_initial_values,\n",
    "                           recon=recon_previous, normalize=normalize_previous,\n",
    "                           remove_global_transform=True, antialiasing_filter=False)\n",
    "\n",
    "        # create dataset:\n",
    "        dataset = a.generate_dataset()\n",
    "\n",
    "        # optimization loop:\n",
    "        ii = 0\n",
    "        losses = list()\n",
    "        for stack_downsamp, rc_downsamp in dataset:\n",
    "            start = time()\n",
    "\n",
    "            loss_i = a.gradient_update(stack_downsamp, rc_downsamp, dither_coords=dither_coords, \n",
    "                                       reg_coefs=reg_coefs, return_loss_only=True)\n",
    "            if type(loss_i) is list:\n",
    "                losses.append([loss.numpy() for loss in loss_i])\n",
    "            else:\n",
    "                losses.append(loss_i.numpy())\n",
    "\n",
    "            if ii % 250 == 0:\n",
    "                print(ii, losses[-1], time() - start)\n",
    "                \n",
    "            if any(np.isnan(loss_i)) or loss_i[0] > 2e3:\n",
    "                raise ValueError('NaN occurred or diverged, restarting ...') \n",
    "\n",
    "            if ii == num_iter:\n",
    "                break\n",
    "            else:\n",
    "                ii += 1\n",
    "\n",
    "        break  # break out of while loop if you successfully complete the optimization loop without errors;\n",
    "\n",
    "    except ValueError:\n",
    "        print('NaN occurred or diverged; restarting the optimization ...')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# when all is done, plot results:\n",
    "loss_i, recon, normalize, error_map, tracked = a.gradient_update(stack_downsamp, rc_downsamp, \n",
    "                                                                 dither_coords=dither_coords, reg_coefs=reg_coefs,\n",
    "                                                                 return_tracked_tensors=True)\n",
    "monitor_progress(recon, error_map, losses, tracked)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tf2_lms",
   "language": "python",
   "name": "tf2_lms"
  },
  "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.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}