<reponame>kevinczhou/mesoscopic-photogrammetry
import numpy as np
import os
import matplotlib.pyplot as plt
import cv2
import tensorflow as tf
from tqdm.notebook import tqdm
import scipy.signal
from tensorflow.python.training.tracking.data_structures import ListWrapper


class mesoSfM:
    def __init__(self, stack, ul_coords, recon_shape, ul_offset=(0, 0), batch_size=None, scale=1, momentum=None,
                 batch_across_images=False, report_error_map=False, restrict_function='clip'):
        # stack: stack of images to be stitched; 3D tensor of shape num_images, num_rows, num_cols, num_channels;
        # ul_coords: array of upper-left coordinates of the images (in pixels); 2D tensor of shape num_images, 2 (row,
        # col; y, x);
        # recon_shape: shape of the final reconstruction/stitched image (row, col, channels);
        # ul_offset: length-2 vector (row, col) specifying a constant offset for all the ul_coords (in pixels);
        # batch_size: if None, then don't use batching;
        # scale: factor between 0 and 1 specifying how much to downsample the reconstruction;
        # momentum: only relevant for batching; specifies how much averaging of previous iterations to use for recon;
        # either both batch_size and momentum must be non-None, or both must be None;
        # batch_across_images: if you're batching (i.e., batch_size is not None), then if true, then batch across the
        # image dimension; otherwise, batch across pixels; (if batching while using unet, then you must batch across
        # image dimension);
        # report_error_map:
        # restrict_function choice decides what to do if a point goes beyond the boundaries; 'clip' or 'mod' or
        # 'mod_with_random_shifts';

        self.tf_dtype = tf.float32
        self.np_dtype = np.float32

        self.stack = np.uint8(stack)  # cast is 8-bit to save memory for the batch generator; cast batch to float32;
        self.num_channels = self.stack.shape[3]  # number of channels; stack must at least have a singleton dim 3;
        self.num_images = self.stack.shape[0]  # number of images in dataset;
        self.ul_coords = np.uint16(ul_coords)  # uint8 is too narrow, because these pixel coordinates are large;
        self.recon_shape_base = recon_shape  # this will be the base recon_shape; the effective recon_shape will depend
        # on the scale factor that the user specifies;
        self.ul_offset = np.array(ul_offset)
        self.scale = scale
        self.momentum = momentum
        self.batch_size = batch_size
        self.batch_across_images = batch_across_images
        self.report_error_map = report_error_map
        self.restrict_function = restrict_function
        self.sig_proj = .42465  # for the intepolation kernel width;
        self.subtract_min_from_height_map = True
        self.optimizer = tf.keras.optimizers.Adam

        # unet parameters if relevant; define these manually if needed;
        self.filters_list = None
        self.skip_list = None
        self.unet_scale = .01  # to scale the output of the unet
        self.output_nonlinearity = 'linear'  # 'linear' or 'leaky_relu'
        self.upsample_method = 'bilinear'  # 'bilinear' or 'nearest'
        self.ckpt = None  # for checkpointing models when using unet;
        self.save_iter = 15  # save model every this many iterations;
        self.recompute_CNN = False  # save memory using tf.recompute_grad;

        self.height_scale_factor = 4000  # to scale the ego_height parameter to avoid dealing with large values;

        # camera parameters for getting absolute scale;
        self.use_absolute_scale_calibration = False
        self.effective_focal_length_mm = 4.3  # effective focal length in mm;
        self.magnification_j = None  # user needs to figure out the magnification of the jth camera view;
        self.j = 0  # and specify which camera view index; by default, assume the first camera;

    def create_variables(self, deformation_model, learning_rates=None, variable_initial_values=None, recon=None,
                         normalize=None, remove_global_transform=False, antialiasing_filter=False,
                         stack_downsample_factor=None, force_ground_surface_up=False):
        # define tf.Variables and optimizers;
        # deformation_model: affects what tf.Variables will be defined and optimized;
        # learning_rates is a dictionary of variable names (strings) and the corresponding learning rate; None means
        # use the default, defined below; can specify a learning rate a negative number, indicating that variable is not
        # optimized; can supply initial variable values as a dictionary; if None, use default;
        # if recon and normalize are supplied, then initialize recon_previous and normalize_previous with these after
        # upsampling; both must be supplied; (this is for momentum/running average calculations); most likely, recon and
        # normalize will be from earlier reconstruction attempts using mesoSfM, but at lower resolution;
        # remove_global_transform: e.g., no global shift or scale;
        # stack_downsample_factor is an integer that downsamples the stack and coordinates; if None, then it will be
        # computed from self.scale;
        # force_ground_surface_up; when using camera model, force the mean surface normal to be [0,0,-1];

        # these must be both be None, or neither None:
        assert (self.batch_size is not None) == (self.momentum is not None)
        assert (recon is None) == (normalize is None)

        # define downsample factor:
        self.recon_shape = np.int32(self.recon_shape_base * self.scale)
        if stack_downsample_factor is None:
            self.downsample = np.int32(1 / self.scale)  # also downsample the images to save computation;
            self.downsample = np.maximum(self.downsample, 1)  # obviously can't downsample with 0;
        else:
            self.downsample = stack_downsample_factor
        self.im_downsampled_shape = np.array([(self.stack.shape[1] - 1) // self.downsample + 1,
                                              (self.stack.shape[2] - 1) // self.downsample + 1])

        # coordinates of the images:
        c = np.arange(self.stack.shape[2], dtype=np.uint16)
        r = np.arange(self.stack.shape[1], dtype=np.uint16)
        r, c = np.meshgrid(r, c, indexing='ij')
        rc_base = np.stack([r, c]).T
        self.rc_base = np.tile(rc_base[None],
                               [self.num_images, 1, 1, 1])  # base coordinates (could be downsampled after applying a
        # scale factor);

        self.deformation_model = deformation_model
        self.remove_global_transform = remove_global_transform
        self.force_ground_surface_up = force_ground_surface_up  # only relevant if using a camera model;

        # list of tf.Variables and optimizers, to be populated by one or more _create_variables_*deformation_model*;
        self.train_var_list = list()
        self.optimizer_list = list()
        self.non_train_list = list()  # list of variables that aren't trained (probably .assign()'d; for checkpoints);
        self.tensors_to_track = dict()  # intermediate tensors to track; have a tf.function return the contents;

        def use_default_for_missing(input_dict, default_dict):
            # to be used directly below; allows for dictionaries in which not all keys are specified; if not specified,
            # then use default_dict's value;

            if input_dict is None:  # if nothing given, then use the default;
                return default_dict
            else:
                for key in default_dict:
                    if key in input_dict:
                        if input_dict[key] is None:  # if the key is present, but None is specified;
                            input_dict[key] = default_dict[key]
                        else:  # i.e., use the value given;
                            pass
                    else:  # if key is not even present;
                        input_dict[key] = default_dict[key]
                return input_dict

        if 'camera_parameters' in deformation_model:
            if 'unet' in deformation_model:
                # user has to manually define these:
                assert self.filters_list is not None
                assert self.skip_list is not None

            default_learning_rates = {'camera_focal_length': 1e-3, 'camera_height': 1e-3, 'ground_surface_normal': 1e-3,
                                      'camera_in_plane_angle': 1e-3, 'rc': 10, 'gain': 1e-3, 'ego_height': 1e-3,
                                      'bias': 1e-3}
            default_variable_initial_values = {'camera_focal_length': np.float32(1),  # initialize same value as ...
                                               'camera_height': 1 * np.ones([self.num_images]),  # ...for height;
                                               'ground_surface_normal': np.concatenate(
                                                   [np.zeros((self.num_images, 2)) + 1e-7,  # to avoid /0;
                                                    -np.ones((self.num_images, 1))], axis=1),
                                               'camera_in_plane_angle': np.zeros([self.num_images]),
                                               'rc': self.ul_coords,
                                               'ego_height': 1e-7 +  # add small value to allow gradients to prop;
                                                         np.zeros([self.num_images,
                                                                   self.im_downsampled_shape[0],
                                                                   self.im_downsampled_shape[1]]),
                                               'gain': np.ones(self.num_images),
                                               'bias': np.zeros(self.num_images)}
            learning_rates = use_default_for_missing(learning_rates, default_learning_rates)
            variable_initial_values = use_default_for_missing(variable_initial_values, default_variable_initial_values)
            self._create_variables_camera_parameters(learning_rates, variable_initial_values)

            # these are for perspective to perspective and perspective to orthographic:
            if 'unet' not in deformation_model:
                # create radial deformation field:
                self._create_variables_perspective_to_orthographic(learning_rates, variable_initial_values)
            elif 'unet' in deformation_model:
                # create neural network layers:
                self._create_variables_height_map_unet(learning_rates, variable_initial_values)

            if self.remove_global_transform:
                # removing the global scaling transform entails adapting the focal length to the mean height;
                assert learning_rates['camera_focal_length'] < 0

            if 'perspective_to_perspective' in deformation_model:
                # these won't be optimized; also these initializing values are very temporary, and should be immediately
                # modified by code later; the reason why they are tf.Variables is so that the user can manually modify
                # them in eager mode with .assign();
                self.reference_camera_height = tf.Variable(1, dtype=self.tf_dtype, name='reference_camera_height')
                self.reference_camera_rc = tf.Variable(np.zeros(2), dtype=self.tf_dtype, name='reference_camera_rc')
                # note that because there's no tilt, reference_camera_height is the same as the camera to vanishing
                # point distance, and the reference_camera_rc is the same as the vanishing point position;
                self.non_train_list.append(self.reference_camera_rc)
                self.non_train_list.append(self.reference_camera_height)
        else:
            raise Exception('invalid deformation model: ' + deformation_model)

        # intensity adjustment (uniform within a given image for now):
        self.gain = tf.Variable(variable_initial_values['gain'], dtype=self.tf_dtype, name='gain')
        self.gain_optim = self.optimizer(learning_rate=learning_rates['gain'])
        self.train_var_list.append(self.gain)
        self.optimizer_list.append(self.gain_optim)

        # intensity bias (also uniform within each image):
        self.bias = tf.Variable(variable_initial_values['bias'], dtype=self.tf_dtype, name='bias')
        self.bias_optim = self.optimizer(learning_rate=learning_rates['bias'])
        self.train_var_list.append(self.bias)
        self.optimizer_list.append(self.bias_optim)

        # barrel or pincushion distortion correction:
        if 'radial_camera_distortion' in learning_rates:
            assert 'camera' in deformation_model  # this is explicitly a camera modeling option;
            if 'radial_camera_distortion' in variable_initial_values:
                camera_distortion_init = variable_initial_values['radial_camera_distortion']
                if np.ndim(camera_distortion_init) == 0:  # user is specifying there to be only one common parameter;
                    camera_distortion_init = np.reshape(camera_distortion_init, (1, 1))  # make sure at least 2 dims;
                elif np.ndim(camera_distortion_init) == 1:
                    # assume this one dimension refers to the camera view, not the polynomial order;
                    camera_distortion_init = np.reshape(camera_distortion_init, (-1, 1))
            else:
                camera_distortion_init = np.zeros((self.num_images, 1))  # second dim: polynomial order;
            num_poly_terms = camera_distortion_init.shape[1]  # number of terms in the polynomial (even only);
            self.radial_powers = (np.arange(num_poly_terms) + 1)  # half of the even powers to raise to;
            self.correct_radial_camera_distortion = camera_distortion_init.shape[0]  # this info used if batching;
            self.radial_camera_distortion = tf.Variable(camera_distortion_init, dtype=self.tf_dtype,
                                                        name='radial_camera_distortion')
            self.radial_camera_distortion_optim = self.optimizer(
                learning_rate=learning_rates['radial_camera_distortion'])
            self.train_var_list.append(self.radial_camera_distortion)
            self.optimizer_list.append(self.radial_camera_distortion_optim)
        else:
            self.correct_radial_camera_distortion = False

        if 'radial_camera_distortion_piecewise_linear' in learning_rates:
            assert 'camera' in deformation_model
            if 'radial_camera_distortion_piecewise_linear' in variable_initial_values:
                camera_distortion_init = variable_initial_values['radial_camera_distortion_piecewise_linear']
                assert np.ndim(camera_distortion_init) == 1  # for now, only allow a common distortion among all cams;
                # length of this vector determines how many discretization levels;
            else:
                camera_distortion_init = np.zeros(50)
            self.num_radial_pixels = len(camera_distortion_init)  # how many discretization levels (nodes);
            self.correct_radial_camera_distortion_piecewise_linear = -1  # used if batching images; this never equals
            # num_images;
            self.radial_camera_distortion_piecewise_linear = tf.Variable(camera_distortion_init, dtype=self.tf_dtype,
                                                                         name='radial_camera_distortion_piecewise_linear')
            self.radial_camera_distortion_piecewise_linear_optim = self.optimizer(
                learning_rate=learning_rates['radial_camera_distortion_piecewise_linear'])
            self.train_var_list.append(self.radial_camera_distortion_piecewise_linear)
            self.optimizer_list.append(self.radial_camera_distortion_piecewise_linear_optim)
        else:
            self.correct_radial_camera_distortion_piecewise_linear = False

        # if camera is not centered; this defines the center of the above two distortions, so center the camera first,
        # apply the distortions, then decenter back;
        if 'camera_distortion_center' in learning_rates:
            assert 'camera' in deformation_model
            if 'camera_distortion_center' in variable_initial_values:
                camera_distortion_center_init = variable_initial_values['camera_distortion_center']
                if np.ndim(camera_distortion_center_init) == 1:  # one common pair of parameters;
                    assert len(camera_distortion_center_init) == 2  # x and y centers;
                    camera_distortion_center_init = tf.reshape(camera_distortion_center_init, (1, -1))
                if np.ndim(camera_distortion_center_init) == 0:
                    raise Exception('must supply two values for x/y centration parameters')
            else:
                camera_distortion_center_init = np.zeros((self.num_images, 2))  # by default, diff pair for each camera;
            self.correct_camera_distortion_center = camera_distortion_center_init.shape[0]  # this info used if
            # batching;
            self.camera_distortion_center = tf.Variable(camera_distortion_center_init, dtype=self.tf_dtype,
                                                    name='camera_distortion_center')
            self.camera_distortion_center_optim = self.optimizer(
                learning_rate=learning_rates['camera_distortion_center'])
            self.train_var_list.append(self.camera_distortion_center)
            self.optimizer_list.append(self.camera_distortion_center_optim)
        else:
            self.correct_camera_distortion_center = False

        # create a list of booleans to accompany self.train_var_list and self.optimizer_list to specify whether to train
        # those variables (as specified by the whether the user-specified learning rates are negative); doing this so
        # that autograph doesn't traverse all branches of the conditionals; if the user ever wants to turn off
        # optimization of a variable mid-optimization, then just do .assign(0) to the learning rate, such that the
        # update is still happening, but the change is 0;
        self.trainable_or_not = list()
        for var in self.train_var_list:
            if type(var) is list:
                # if the variable is a list of variables, then this should be for the unet; modify here if there are
                # other scenarios;
                assert 'unet' in self.deformation_model
                name = 'ego_height'
            else:
                name = var.name[:-2]
            flag = learning_rates[name] > 0
            self.trainable_or_not.append(flag)

        # downsample rc coordinates and stack:
        rc = np.transpose(self.rc_base[:, ::self.downsample, ::self.downsample, :], (0, 2, 1, 3))
        if antialiasing_filter:
            downsample = int(self.downsample)  # cv2 doesn't like numpy values?
            if downsample == 1:
                print('warning: antialiasing filter is applied even though there is no downsampling')
            ksize = int(downsample * 2.5)
            if ksize % 2 == 0:  # must be odd;
                ksize += 1
            stack_downsamp = np.stack(
                [cv2.GaussianBlur(im, (ksize, ksize), downsample, downsample)
                 [::downsample, ::downsample] for im in self.stack])
        else:
            stack_downsamp = self.stack[:, ::self.downsample, ::self.downsample, :]

        rc_downsamp = np.reshape(rc, (self.rc_base.shape[0], -1, self.rc_base.shape[-1]))  # flatten spatial dims;
        stack_downsamp = np.reshape(stack_downsamp,
                                    [self.num_images, -1, self.num_channels])  # flatten spatial dims;
        self.rc_downsamp = rc_downsamp
        self.stack_downsamp = stack_downsamp

        # create variables relevant for batching (or set them to None if not batching):
        if self.momentum is not None:
            if 'camera_parameters_perspective_' in deformation_model:
                # we're going to give the coregistered height map a ride;
                num_channels = self.num_channels + 1
            else:
                num_channels = self.num_channels

            if recon is None:  # if none supplied, initialize with 0s;
                recon_previous = np.zeros([self.recon_shape[0], self.recon_shape[1], num_channels])
                normalize_previous = np.zeros(self.recon_shape)
            else:  # otherwise, upsample to the current shape;
                recon_previous = cv2.resize(np.nan_to_num(recon), tuple(self.recon_shape[::-1]))
                normalize_previous = cv2.resize(np.nan_to_num(normalize), tuple(self.recon_shape[::-1]))
                if num_channels == 1:
                    # cv2 seems to squeeze singleton channels dimensions, so if it's singleton, add it back:
                    recon_previous = recon_previous[:, :, None]
                if recon_previous.shape[-1] != num_channels:
                    # if you would like to use the previous RGB image as initialization, but your previous run didn't
                    # estimate a height map;
                    assert 'camera_parameters_perspective_' in deformation_model
                    # add empty height channel:
                    recon_previous = np.concatenate([recon_previous, np.zeros_like(recon_previous[:, :, 0:1])], axis=-1)

            # initialize first recon and normalize tensors for momentum; use the scaled recon shape, not the base shape;
            self.recon_previous = tf.Variable(recon_previous, dtype=self.tf_dtype, trainable=False)
            self.non_train_list.append(self.recon_previous)
        else:
            self.recon_previous = None
            self.normalize_previous = None

    def _create_variables_perspective_to_orthographic(self, learning_rates, variable_initial_values):
        # radially inwardly pointing vector magnitudes, where the larger the magnitude, the taller the object;

        ego_height = variable_initial_values['ego_height']
        # make sure the first dimensions match (number of images in stack):
        assert ego_height.shape[0] == self.num_images
        if ego_height.shape[1:] != tuple(self.im_downsampled_shape):
            # presumably you've initialized this with the results from another optimization at a different scale;
            # thus, resize to match the current scale:
            if type(ego_height) != np.ndarray:
                # convert from tf to np if you need to:
                ego_height = ego_height.numpy()
            ego_height = np.stack([cv2.resize(im, tuple(self.im_downsampled_shape[::-1])) for im in ego_height])

        self.ego_height = tf.Variable(ego_height, dtype=self.tf_dtype, name='ego_height')
        self.ego_height_optim = self.optimizer(learning_rate=learning_rates['ego_height'])
        self.train_var_list.append(self.ego_height)
        self.optimizer_list.append(self.ego_height_optim)

    def _create_variables_height_map_unet(self, learning_rates, variable_initial_values):

        self.network = unet(self.filters_list, self.skip_list, output_nonlinearity=self.output_nonlinearity,
                            upsample_method=self.upsample_method)
        # run the network once so that we can access network.trainable_variables
        self.network(tf.zeros([1, 2 ** len(self.filters_list), 2 ** len(self.filters_list), self.num_channels],
                              dtype=self.tf_dtype))

        self.train_var_list.append(self.network.trainable_variables)
        self.optimizer_list.append(self.optimizer(learning_rate=learning_rates['ego_height']))

        if self.recompute_CNN:
            self.network = tf.recompute_grad(self.network)

        # get padded shape that the network likes;
        self.padded_shape = [get_compatible_size(dim, len(self.filters_list)) for dim in self.im_downsampled_shape]
        pad_r = self.padded_shape[0] - self.im_downsampled_shape[0]
        pad_c = self.padded_shape[1] - self.im_downsampled_shape[1]
        pad_top = pad_r // 2
        pad_bottom = int(tf.math.ceil(pad_r / 2))
        pad_left = pad_c // 2
        pad_right = int(tf.math.ceil(pad_c / 2))
        pad_specs = ((pad_top, pad_bottom), (pad_left, pad_right))
        self.pad_layer = tf.keras.layers.ZeroPadding2D(pad_specs)
        self.depad_layer = tf.keras.layers.Cropping2D(pad_specs)

    def _create_variables_camera_parameters(self, learning_rates, variable_initial_values):
        # coordinates are in camera space;

        self.camera_focal_length = tf.Variable(variable_initial_values['camera_focal_length'], dtype=self.tf_dtype,
                                               name='camera_focal_length')
        self.camera_focal_length_optim = self.optimizer(learning_rate=learning_rates['camera_focal_length'])

        self.camera_height = tf.Variable(variable_initial_values['camera_height'], dtype=self.tf_dtype,
                                         name='camera_height')  # height from the camera perspective;
        self.camera_height_optim = self.optimizer(learning_rate=learning_rates['camera_height'])

        self.ground_surface_normal = tf.Variable(variable_initial_values['ground_surface_normal'], dtype=self.tf_dtype,
                                                 name='ground_surface_normal')
        self.ground_surface_normal_optim = self.optimizer(
            learning_rate=learning_rates['ground_surface_normal'])

        self.camera_in_plane_angle = tf.Variable(variable_initial_values['camera_in_plane_angle'],
                                                 dtype=self.tf_dtype, name='camera_in_plane_angle')
        self.camera_in_plane_angle_optim = self.optimizer(
            learning_rate=learning_rates['camera_in_plane_angle'])

        self.rc_ul_per_im = tf.Variable(variable_initial_values['rc'], dtype=self.tf_dtype, name='rc')
        self.rc_ul_per_im_optim = self.optimizer(learning_rate=learning_rates['rc'])

        self.train_var_list.append(self.camera_focal_length)
        self.optimizer_list.append(self.camera_focal_length_optim)
        self.train_var_list.append(self.camera_height)
        self.optimizer_list.append(self.camera_height_optim)
        self.train_var_list.append(self.ground_surface_normal)
        self.optimizer_list.append(self.ground_surface_normal_optim)
        self.train_var_list.append(self.camera_in_plane_angle)
        self.optimizer_list.append(self.camera_in_plane_angle_optim)
        self.train_var_list.append(self.rc_ul_per_im)
        self.optimizer_list.append(self.rc_ul_per_im_optim)

    def generate_dataset(self):
        # user calls this function to get a dataset to iterate over; if not using batching, then just return a tuple or
        # list of length 1 (i.e., the whole dataset is one batch);

        if self.batch_size is not None:
            if self.batch_across_images:
                # sample a subset of the images, and keep track of the indices downsampled so that you can gather the
                # corresponding variables;
                tensor_slices = (self.stack_downsamp,
                                 (self.rc_downsamp, np.arange(self.num_images, dtype=np.int32)))
                dataset = (tf.data.Dataset.from_tensor_slices(tensor_slices).shuffle(self.num_images)
                           .batch(self.batch_size, drop_remainder=True).repeat(None).prefetch(1))
                return dataset
            else:
                # transpose to batch along space, not image number;
                rc_downsamp_T = np.transpose(self.rc_downsamp, (1, 0, 2))
                stack_downsamp_T = np.transpose(self.stack_downsamp, (1, 0, 2))
                if 'camera_parameters_perspective_' in self.deformation_model:
                    # need to also get coordinates of the spatial positions to index into pixel-wise deformation fields:
                    tensor_slices = (stack_downsamp_T,
                                     (rc_downsamp_T, np.arange(np.prod(self.im_downsampled_shape), dtype=np.int32)))
                else:
                    tensor_slices = (stack_downsamp_T, rc_downsamp_T)
                dataset = (tf.data.Dataset.from_tensor_slices(tensor_slices)
                           .shuffle(len(rc_downsamp_T)).batch(self.batch_size).repeat(None).prefetch(1))
                return dataset
        else:  # basically a 1-batch dataset;
            return self.stack_downsamp, self.rc_downsamp

    def _warp_camera_parameters(self, rc_downsamp, use_radial_deformation, p2p_warp_mode=None,
                                inds_downsamp=None, stack_downsamp=None):
        # shape of rc_downsamp: num_images, _, 2;
        # use_radial_deformation is a boolean flag specifying whether to do the per-pixel radial deformation fields to
        # warp perspective to othographic OR perspective to perspective; if the latter, then p2p_warp_mode specifies
        # how to specify the perspective reference to warp to; the options are 'mean', 'random', 'fixed', and None,
        # where None means you're using perspective-to-orthographpic warping and not perspective to perspective;
        # inds_downsamp is passed if batching and using radial deformations;
        # stack_downsamp is only needed if using a unet;

        if p2p_warp_mode is None:
            assert 'perspective_to_orthographic' in self.deformation_model
        else:
            assert 'perspective_to_perspective' in self.deformation_model

        if self.batch_across_images and self.batch_size is not None:
            # in generate_recon, we defined the batch and non-batch versions;

            rc_ul_per_im = self.rc_ul_per_im_batch
            gain = self.gain_batch
            bias = self.bias_batch
            camera_height = self.camera_height_batch
            ground_surface_normal = self.ground_surface_normal_batch
            camera_in_plane_angle = self.camera_in_plane_angle_batch
            if 'unet' not in self.deformation_model:
                ego_height = self.ego_height_batch
                self.ego_height_to_regularize = self.ego_height_batch
            else:
                pass
            if self.correct_radial_camera_distortion:
                if self.correct_radial_camera_distortion == self.num_images:
                    radial_camera_distortion = self.radial_camera_distortion_batch
                else:
                    radial_camera_distortion = self.radial_camera_distortion
            if self.correct_radial_camera_distortion_piecewise_linear:
                if self.correct_radial_camera_distortion_piecewise_linear == self.num_images:
                    radial_camera_distortion_piecewise_linear = self.radial_camera_distortion_piecewise_linear_batch
                else:
                    radial_camera_distortion_piecewise_linear = self.radial_camera_distortion_piecewise_linear
            if self.correct_camera_distortion_center:
                if self.correct_camera_distortion_center == self.num_images:
                    camera_distortion_center = self.camera_distortion_center_batch
                else:
                    camera_distortion_center = self.camera_distortion_center
            num_images = self.batch_size  # for reshaping below;
            camera_focal_length = self.camera_focal_length  # not batched;
        else:
            rc_ul_per_im = self.rc_ul_per_im
            gain = self.gain
            bias = self.bias
            camera_height = self.camera_height
            ground_surface_normal = self.ground_surface_normal
            camera_in_plane_angle = self.camera_in_plane_angle
            if 'unet' not in self.deformation_model:
                ego_height = self.ego_height
                self.ego_height_to_regularize = self.ego_height
            if self.correct_radial_camera_distortion:
                radial_camera_distortion = self.radial_camera_distortion
            if self.correct_radial_camera_distortion_piecewise_linear:
                radial_camera_distortion_piecewise_linear = self.radial_camera_distortion_piecewise_linear
            if self.correct_camera_distortion_center:
                camera_distortion_center = self.camera_distortion_center

            num_images = self.num_images  # for reshaping below;
            camera_focal_length = self.camera_focal_length

        if self.remove_global_transform:
            # don't use self.camera_focal_length; set to geometric mean of the camera heights;
            # also, use the tf.Variable version always (not the batch version);
            camera_focal_length = tf.reduce_prod(self.camera_height, axis=0, keepdims=False) ** (1 / self.num_images)
            self.camera_focal_length.assign(camera_focal_length)
            camera_in_plane_angle = camera_in_plane_angle - tf.reduce_mean(self.camera_in_plane_angle)
            rc_ul_per_im = rc_ul_per_im - tf.reduce_mean(self.rc_ul_per_im, axis=0, keepdims=True)

        im_dims = np.array(self.stack.shape)[1:3]  # for normalization of image coordinates;
        max_dim = np.max(im_dims)  # to keep isotropic;
        camera_yx = (rc_downsamp - .5 * im_dims[None, None, :]) / max_dim  # normalize to -.5 to .5;

        if self.correct_camera_distortion_center:
            camera_yx -= camera_distortion_center[:, None, :]

        if self.correct_radial_camera_distortion:
            camera_r2 = camera_yx[:, :, 0] ** 2 + camera_yx[:, :, 1] ** 2  # radial distance squared;
            # dims^: camera, pixels
            camera_r2 *= 2  # make it go from -1 to 1 rather than -.5 to .5;

            if self.correct_radial_camera_distortion:
                # even polynomial to account for distortion:
                camera_even_poly = tf.math.pow(camera_r2[:, :, None], self.radial_powers[None, None, :])
                # dims^: camera, pixels, power
                camera_even_poly = tf.reduce_sum(camera_even_poly * radial_camera_distortion[:, None, :], 2)
                # dims^: camera, pixels
                radial_correction_factor = 1 + camera_even_poly[:, :, None]
                self.tensors_to_track['camera_distortion_radial'] = radial_correction_factor
            else:
                radial_correction_factor = 1
            camera_yx = camera_yx * radial_correction_factor

        if self.correct_radial_camera_distortion_piecewise_linear:
            camera_r = tf.sqrt(camera_yx[:, :, 0] ** 2 + camera_yx[:, :, 1] ** 2)  # radial distance;
            # dims^: camera, pixels; these radial coordinates should be between 0 and .5*sqrt(2), but could go higher
            # if th center moves); thus to be safe, just multiply by num_radial_pixels;
            r_scale = camera_r * self.num_radial_pixels
            # find nearest pixels and distances thereto:
            r_floor = tf.floor(r_scale)
            r_ceil = tf.minimum(r_floor + 1, self.num_radial_pixels - 1)  # to prevent out of range indexing;
            r_middle = r_scale - r_floor
            r_floor = tf.cast(r_floor, dtype=tf.int32)
            r_ceil = tf.cast(r_ceil, dtype=tf.int32)
            distortion = 1 + radial_camera_distortion_piecewise_linear
            distortion /= tf.reduce_max(distortion)  # to prevent global expansion;
            correction_factor_floor = tf.gather(distortion, r_floor)
            correction_factor_ceil = tf.gather(distortion, r_ceil)
            # bilinear interpolation:
            correction_factor = correction_factor_ceil * r_middle + correction_factor_floor * (1 - r_middle)
            camera_yx *= correction_factor[:, :, None]
            self.tensors_to_track['camera_distortion_piecewise_linear'] = correction_factor

        if self.correct_camera_distortion_center:
            camera_yx += camera_distortion_center[:, None, :]

        # in-plane rotation:
        cos = tf.cos(camera_in_plane_angle)
        sin = tf.sin(camera_in_plane_angle)
        rotmat_xy = tf.stack([[cos, sin], [-sin, cos]])
        camera_yx = tf.einsum('cri,ijc->crj', camera_yx, rotmat_xy)

        n_ground, _ = tf.linalg.normalize(ground_surface_normal, axis=1)  # normalize to unit mag;
        # shape^: num_images, 3

        # projecting to object space (computed analytically and taylor-expanded);
        nx = n_ground[:, 0][:, None, None]
        ny = n_ground[:, 1][:, None, None]
        nx2 = nx ** 2
        ny2 = ny ** 2
        x = camera_yx[:, :, 1][:, :, None]
        y = camera_yx[:, :, 0][:, :, None]
        h = camera_height[:, None, None]
        f = camera_focal_length
        # using a taylor expansion:
        flat_x = (h * x / f +
                  h * x * (nx * x + ny * y) / f ** 2 +
                  h * (f ** 2 * nx2 * x + 2 * nx2 * x ** 3 + f ** 2 * nx * ny * y + 4 * nx * ny * x ** 2 * y +
                       2 * ny2 * x * y ** 2) / 2 / f ** 3)
        flat_y = (h * y / f +
                  h * y * (nx * x + ny * y) / f ** 2 +
                  h * (f ** 2 * ny2 * y + 2 * ny2 * y ** 3 + f ** 2 * nx * ny * x + 4 * nx * ny * y ** 2 * x +
                       2 * nx2 * y * x ** 2) / 2 / f ** 3)
        flat_xy = tf.concat([flat_x, flat_y], axis=2)

        n_dot_r = n_ground[:, 2] * camera_height  # shape: num_images; dot product between a point on the
        # ground, r (0, 0, camera_height); (this is needed below);
        self.flat_xy = flat_xy

        if use_radial_deformation:
            # compute the vanishing point, which will be used if you use the projective to orthographic mode;
            vanish_xyz = n_dot_r[:, None] * n_ground  # shape: num_images, 3;
            camera_to_vanish_point_xyz = tf.norm(vanish_xyz, axis=1)  # distance from camera to the ground; the actual
            # height of the camera;

            # projection to object space simplifies to this:
            vanish_xy = -camera_height[:, None] * n_ground[:, :2]

            # vanishing point in camera plane:
            vanish_camera_xyz = n_ground * camera_focal_length / n_ground[:, 2:]  # follow surface normal to camera
            # plane;
            vanish_camera_xy = vanish_camera_xyz[:, 0:2]  # don't need z, it's just the focal length;
            # account for in-plane camera rotation:
            vanish_camera_xy = tf.einsum('ci,ijc->cj', vanish_camera_xy, rotmat_xy)
            self.vanish_xy = vanish_xy

        # convert back to row-column:
        flat_rc = (flat_xy[:, :, ::-1] * max_dim + .5 * im_dims[None, None, :])
        if use_radial_deformation:
            vanish_rc = (vanish_xy[:, ::-1] * max_dim + .5 * im_dims[None, :])
            vanish_camera_rc = (vanish_camera_xy[:, ::-1] * max_dim + .5 * im_dims[None, :])
            self.camera_to_vanish_point_rc = camera_to_vanish_point_xyz * max_dim  # convert from xy units to rc units;
            self.tensors_to_track['camera_to_vanish_point_xyz'] = self.camera_to_vanish_point_rc

        # add translations (same as for the homography and affine implementations):
        # these translations don't affect camera_to_vanish_point;
        rc_warp = rc_ul_per_im[:, None, :] + flat_rc
        if use_radial_deformation:
            vanish_warp = rc_ul_per_im + vanish_rc
        if self.restrict_function == 'mod_with_random_shifts':
            # to discourage registration with overlapped regions;
            random_shift = tf.random.uniform(shape=(1, 1, 2), minval=0, maxval=self.recon_shape.max() / self.scale)
            rc_warp += random_shift
            if use_radial_deformation:
                vanish_warp += random_shift[0]
        else:
            rc_warp += self.ul_offset[None, None, :]
            if use_radial_deformation:
                vanish_warp += self.ul_offset[None, :]

        if 'unet' in self.deformation_model:
            # generate self.ego_height; doesn't matter if batching or not because it's generated from the image batch;
            unet_input = tf.reshape(stack_downsamp, [num_images, self.im_downsampled_shape[0],
                                                     self.im_downsampled_shape[1], self.num_channels])

            unet_input = self.pad_layer(unet_input)
            unet_output = self.network(unet_input)
            ego_height = tf.reduce_mean(self.depad_layer(unet_output), [-1])  # remove last dimension;
            ego_height *= self.unet_scale
            self.ego_height_to_regularize = ego_height  # need this for regularization;
            self.ego_height = ego_height

        # multiplicative version:
        if use_radial_deformation:
            self.tensors_to_track['vanish_warp'] = vanish_warp
            self.tensors_to_track['vanish_camera'] = vanish_camera_rc
            if p2p_warp_mode is None:  # using perspective-to-orthographic
                ego_height_flat = tf.reshape(ego_height, [num_images, -1])  # flatten spatial dimensions;
                if self.subtract_min_from_height_map:
                    ego_height_flat -= tf.reduce_min(ego_height_flat)
                if inds_downsamp is not None:
                    assert self.batch_size is not None  # this should never be raised, but just in case;
                    ego_height = tf.gather(ego_height_flat, inds_downsamp, axis=1)  # batch along pixels;
                else:
                    ego_height = ego_height_flat

                if self.use_absolute_scale_calibration:

                    H = self.camera_to_vanish_point_rc[:, None]

                    if self.batch_size is not None:
                        # need to compute the self.j'th camera-to-vanish-point height in case you're batching, which means
                        # that the self.j'th entry may not be computed;
                        n_ground, _ = tf.linalg.normalize(self.ground_surface_normal[self.j], axis=0)
                        n_dot_r = n_ground[2] * self.camera_height[self.j]
                        vanish_xyz = n_dot_r * n_ground  # shape: 3;
                        camera_to_vanish_point_xyz = tf.norm(vanish_xyz, axis=0)
                        H_j = camera_to_vanish_point_xyz * max_dim
                    else:
                        H_j = H[self.j]

                    r = rc_warp - vanish_warp[:, None, :]  # lateral distance to vanishing point;
                    M_j = self.magnification_j
                    f_eff = self.effective_focal_length_mm
                    self.another_height_scale_factor = f_eff * (1 + 1 / M_j)  # scale ego_height again to make this
                    # case similar to the other case (self.use_absolute_scale_calibration);
                    h = ego_height * self.another_height_scale_factor

                    delta_radial = h / f_eff / (1 + 1 / M_j * H / H_j)
                    rc_warp = r * (1 - delta_radial[:, :, None]) + vanish_warp[:, None, :]

                    ego_height *= self.height_scale_factor  # to keep consistent with regularization coefficients;
                    # note that you have to divide by height_scale_factor because the height map is scaled by this, but
                    # you have to divide by another_scale_factor because the multiplication above allows ego_height to
                    # shrink;
                else:
                    ego_height *= self.height_scale_factor  # denominator is large in next line, so multiply by a large
                    # value to allow self.ego_height to take on smaller values;

                    delta_radial = ego_height / self.camera_to_vanish_point_rc[:, None]
                    rc_warp = ((rc_warp - vanish_warp[:, None, :]) * (1 - delta_radial[:, :, None]) +
                               vanish_warp[:, None, :])
                self.ego_height_for_concat = ego_height  # for concatenating with self.im below;
            else:  # perspective-to-perspective warping
                # first, need to define the reference camera view:
                if p2p_warp_mode == 'mean':
                    self.reference_camera_height.assign(tf.reduce_mean(self.camera_to_vanish_point_rc))
                    self.reference_camera_rc.assign(tf.reduce_mean(vanish_warp, axis=0))
                elif p2p_warp_mode == 'random':
                    height_min = tf.reduce_min(self.camera_to_vanish_point_rc)
                    height_max = tf.reduce_max(self.camera_to_vanish_point_rc)
                    self.reference_camera_height.assign(tf.random.uniform((), height_min, height_max))
                    rc_min = tf.reduce_min(vanish_warp, axis=0)
                    rc_max = tf.reduce_max(vanish_warp, axis=0)
                    self.reference_camera_rc.assign(tf.random.uniform((2,), rc_min, rc_max))
                elif p2p_warp_mode == 'fixed':
                    pass  # do nothing, accept current values;
                elif p2p_warp_mode == 'random_choice':
                    # pick one of the camera view among the existing;
                    random_choice = tf.random.uniform((), 0, self.num_images, dtype=tf.int32)
                    self.reference_camera_height.assign(tf.gather(self.camera_to_vanish_point_rc, random_choice))
                    self.reference_camera_rc.assign(tf.gather(vanish_warp, random_choice))
                else:
                    raise Exception('invalid perspective-to-perspective warp mode passed to gradient_update: '
                                    + p2p_warp_mode)
                # vector deformation field to warp to the reference perspective:
                h = tf.reshape(ego_height, [num_images, -1, 1]) * self.height_scale_factor
                H_r = self.reference_camera_height
                H = self.camera_to_vanish_point_rc[:, None, None]
                R_r = self.reference_camera_rc[None, None, :]
                R = vanish_warp[:, None, :]
                r = rc_warp - R  # position vectors relative to each camera's vanishing point;
                p2p_warp = h / (H_r - h) * (R - R_r) + r * h * (H - H_r) / H / (H_r - h)  # the magic equation;
                rc_warp += p2p_warp  # shape: num_images, flattened spatial, 2

                self.ego_height_for_concat = h  # for concatenating with self.im below;

                if self.use_absolute_scale_calibration:
                    raise Exception('not yet implemented for perspective-to-perspective')

        rc_warp = tf.reshape(rc_warp, [-1, 2]) * self.scale  # flatten

        return rc_warp

    def _generate_recon(self, stack_downsamp, rc_downsamp, dither_coords, p2p_warp_mode=None, assign_update_recon=True):
        # backprojects all the images into the reconstruction, with the specified scale;
        # if batching, the gradient_update function will update the reconstruction with a running average;
        # if batching, this function should not be called by the user, as it will continually update the recon with the
        # same batch; if not batching, then this generates the full reconstruction;
        # p2p_warp_mode: if using perspective-to-perspective warping; can be None, 'mean', 'random', or 'fixed';
        # assign_update_recon: only relevant if using batching; controls whether to use the .assign() mechanism to
        # update the reconstruction (specified via update_gradient option in the gradient_update function);

        if self.batch_size is not None:
            if self.batch_across_images:
                # distinguish inds_downsamp and inds_image_downsamp, where the former is for pixel-level batching while
                # the latter is for image-level batching;
                rc_downsamp, inds_image_downsamp = rc_downsamp  # unpack

                # now, for all variables whose first dimension corresponds to the image dimension, gather:
                self.rc_ul_per_im_batch = tf.gather(self.rc_ul_per_im, inds_image_downsamp, axis=0)
                self.gain_batch = tf.gather(self.gain, inds_image_downsamp, axis=0)
                self.bias_batch = tf.gather(self.bias, inds_image_downsamp, axis=0)
                # these are used below:
                gain = self.gain_batch
                bias = self.bias_batch

                if 'camera_parameters' in self.deformation_model:
                    self.camera_height_batch = tf.gather(self.camera_height, inds_image_downsamp, axis=0)
                    self.ground_surface_normal_batch = tf.gather(self.ground_surface_normal,
                                                                 inds_image_downsamp, axis=0)
                    self.camera_in_plane_angle_batch = tf.gather(self.camera_in_plane_angle,
                                                                 inds_image_downsamp, axis=0)
                    if 'unet' not in self.deformation_model:
                        # if using unet, then ego_height will already be gathered as it is generated by the unet;
                        self.ego_height_batch = tf.gather(self.ego_height, inds_image_downsamp, axis=0)
                    # the following self.correct__ variables serve two purposes: 1) to signify whether they are being
                    # used, and 2) specify length of first dimension of the corresponding distortion variable to decide
                    # whether we need to use tf.gather;
                    if self.correct_radial_camera_distortion == self.num_images:
                        self.radial_camera_distortion_batch = tf.gather(self.radial_camera_distortion,
                                                                  inds_image_downsamp, axis=0)
                    if self.correct_radial_camera_distortion_piecewise_linear == self.num_images:
                        self.radial_camera_distortion_piecewise_linear_batch = tf.gather(
                            self.radial_camera_distortion_piecewise_linear, inds_image_downsamp, axis=0)
                    if self.correct_camera_distortion_center == self.num_images:
                        self.camera_distortion_center_batch = tf.gather(self.camera_distortion_center,
                                                                        inds_image_downsamp, axis=0)
                else:
                    raise Exception('image-level batching not yet implemented for a non-camera model')

                inds_downsamp = None
            else:
                # if batching, then stack_downsamp and rc_downsamp are transposed and need to be untransposed;

                if 'camera_parameters' in self.deformation_model:
                    # also need the indices of the pixels chosen, because radial deformations are pixel-wise;
                    # make sure you package these together into a tuple in the script;
                    rc_downsamp, inds_downsamp = rc_downsamp
                else:
                    inds_downsamp = None

                stack_downsamp = tf.transpose(stack_downsamp, (1, 0, 2))
                rc_downsamp = tf.transpose(rc_downsamp, (1, 0, 2))

                # these are used below:
                gain = self.gain
                bias = self.bias
        else:
            inds_downsamp = None
            # these are used below:
            gain = self.gain
            bias = self.bias
        # to save CPU memory, the dataset and coordinates are stored as uint8 and uint16, respectively; thus, cast to
        # float here;
        stack_downsamp = tf.cast(stack_downsamp, self.tf_dtype)
        rc_downsamp = tf.cast(rc_downsamp, self.tf_dtype)

        # function that restricts coordinates to the grid (store as self.variable so that error_map can use it):
        if self.restrict_function == 'clip':
            self.restrict = lambda x: tf.clip_by_value(x, tf.zeros_like(x), self.recon_shape[None] - 1)
        elif 'mod' in self.restrict_function:  # 'mod' or 'mod_with_random_shifts';
            self.restrict = lambda x: tf.math.floormod(x, self.recon_shape[None])
        else:
            raise Exception('invalid restrict_function')

        # apply gain:
        gain_norm = gain / tf.reduce_mean(gain)  # normalize so that there's no global gain;
        im = stack_downsamp * gain_norm[:, None, None] + bias[:, None, None]

        self.im = tf.reshape(im, (-1, self.num_channels))  # flatten all but channels;

        # warped coordinates:
        if self.deformation_model == 'camera_parameters':
            self.rc_warp = self._warp_camera_parameters(rc_downsamp, use_radial_deformation=False)
        elif self.deformation_model == 'camera_parameters_perspective_to_orthographic':
            self.rc_warp = self._warp_camera_parameters(rc_downsamp, inds_downsamp=inds_downsamp,
                                                        use_radial_deformation=True)
        elif self.deformation_model == 'camera_parameters_perspective_to_orthographic_unet':
            self.rc_warp = self._warp_camera_parameters(rc_downsamp, inds_downsamp=inds_downsamp,
                                                        use_radial_deformation=True, stack_downsamp=stack_downsamp)
        elif self.deformation_model == 'camera_parameters_perspective_to_perspective':
            self.rc_warp = self._warp_camera_parameters(rc_downsamp, p2p_warp_mode=p2p_warp_mode,
                                                        inds_downsamp=inds_downsamp,
                                                        use_radial_deformation=True)
        elif self.deformation_model == 'camera_parameters_perspective_to_perspective_unet':
            self.rc_warp = self._warp_camera_parameters(rc_downsamp, p2p_warp_mode=p2p_warp_mode,
                                                        inds_downsamp=inds_downsamp,
                                                        use_radial_deformation=True, stack_downsamp=stack_downsamp)
        else:
            raise Exception('invalid deformation model: ' + self.deformation_model)

        if 'camera_parameters_perspective_to_' in self.deformation_model:
            # adding the height map as a channel to the reconstruction, so first augment self.im with self.ego_height:
            self.im = tf.concat([self.im, tf.reshape(self.ego_height_for_concat, [-1])[:, None]], axis=1)

            self.num_channels_recon = self.num_channels + 1  # for the recon, need one more channel;
        else:
            self.num_channels_recon = self.num_channels

        #
        if dither_coords:
            self.rc_warp += tf.random.uniform([1, 2], -1, 1, dtype=self.tf_dtype)
            if self.batch_size is not None:
                print('Minor warning: using a running average for the recon while dithering coordinates')

        # neighboring pixels:
        rc_floor = tf.floor(self.rc_warp)
        rc_ceil = rc_floor + 1

        # distance to neighboring pixels:
        frc = self.rc_warp - rc_floor
        crc = rc_ceil - self.rc_warp

        # cast
        rc_floor = tf.cast(rc_floor, tf.int32)
        rc_ceil = tf.cast(rc_ceil, tf.int32)

        self.rc_ff = self.restrict(rc_floor)
        self.rc_cc = self.restrict(rc_ceil)
        self.rc_cf = self.restrict(tf.stack([rc_ceil[:, 0], rc_floor[:, 1]], 1))
        self.rc_fc = self.restrict(tf.stack([rc_floor[:, 0], rc_ceil[:, 1]], 1))

        # sig_proj = .42465  # chosen so that if point is exactly in between
        # ...two pixels, .5 weight is assigned to each pixel
        self.frc = tf.exp(-frc ** 2 / 2. / self.sig_proj ** 2)
        self.crc = tf.exp(-crc ** 2 / 2. / self.sig_proj ** 2)

        # augmented coordinates:
        rc_4 = tf.concat([self.rc_ff, self.rc_cc, self.rc_cf, self.rc_fc], 0)

        # interpolated:
        im_4 = tf.concat([self.im * self.frc[:, 0, None] * self.frc[:, 1, None],
                          self.im * self.crc[:, 0, None] * self.crc[:, 1, None],
                          self.im * self.crc[:, 0, None] * self.frc[:, 1, None],
                          self.im * self.frc[:, 0, None] * self.crc[:, 1, None]], 0)
        w_4 = tf.concat([self.frc[:, 0] * self.frc[:, 1],
                         self.crc[:, 0] * self.crc[:, 1],
                         self.crc[:, 0] * self.frc[:, 1],
                         self.frc[:, 0] * self.crc[:, 1]], 0)

        if self.momentum is not None:
            # update with moving average:
            self.im_4_previous = tf.gather_nd(self.recon_previous,
                                              rc_4) * w_4[:, None]  # with appropriate weighting by w_4;
            self.im_4_updated = (im_4 * self.momentum + self.im_4_previous * (1 - self.momentum))
            normalize = tf.scatter_nd(rc_4, w_4, self.recon_shape)
            self.norm_updated_regathered = tf.gather_nd(normalize, rc_4)
            self.im_4_updated_norm = self.im_4_updated / self.norm_updated_regathered[:, None]  # pre-normalize;
            # since tensor_scatter_nd_update doesn't accumulate values, but tensor_scatter_nd_add does, first zero
            # out the regions to be updated and then just add them:
            recon_zeroed = tf.tensor_scatter_nd_update(self.recon_previous, rc_4,
                                                       tf.zeros_like(self.im_4_updated_norm))
            self.recon = tf.tensor_scatter_nd_add(recon_zeroed, rc_4, self.im_4_updated_norm)
            if assign_update_recon:
                with tf.device('/CPU:0'):
                    self.recon_previous.assign(self.recon)

            self.normalize = None  # normalize not needed; in fact, normalize_previous also not needed;

        else:
            self.normalize = tf.scatter_nd(rc_4, w_4, self.recon_shape)
            self.recon = tf.scatter_nd(rc_4, im_4, [self.recon_shape[0], self.recon_shape[1], self.num_channels_recon])
            self.recon = tf.math.divide_no_nan(self.recon, self.normalize[:, :, None])  # creates recon H by W by C;

        if 'camera_parameters_perspective_to' in self.deformation_model:
            self.height_map = self.recon[:, :, -1]
            if self.use_absolute_scale_calibration:
                # divide out the scale factors to get the true height in mm:
                self.tensors_to_track['height_map'] = self.height_map / (self.height_scale_factor /
                                                                         self.another_height_scale_factor)
            else:
                self.tensors_to_track['height_map'] = self.height_map

    def _forward_prediction(self):
        # given the reconstruction, generate forward prediction;

        # forward model:
        ff = tf.gather_nd(self.recon, self.rc_ff)
        cc = tf.gather_nd(self.recon, self.rc_cc)
        cf = tf.gather_nd(self.recon, self.rc_cf)
        fc = tf.gather_nd(self.recon, self.rc_fc)

        self.forward = (ff * self.frc[:, 0, None] * self.frc[:, 1, None] +
                        cc * self.crc[:, 0, None] * self.crc[:, 1, None] +
                        cf * self.crc[:, 0, None] * self.frc[:, 1, None] +
                        fc * self.frc[:, 0, None] * self.crc[:, 1, None])

        self.forward /= ((self.frc[:, 0, None] * self.frc[:, 1, None]) +
                         (self.crc[:, 0, None] * self.crc[:, 1, None]) +
                         (self.crc[:, 0, None] * self.frc[:, 1, None]) +
                         (self.frc[:, 0, None] * self.crc[:, 1, None]))

        if 'camera_parameters_perspective' in self.deformation_model:
            # split off the last dimension, the height dimension, to compute the height map MSE:
            self.forward_height = self.forward[:, -1]
            self.error_height = self.forward_height - self.im[:, -1]  # save this for computing error map;
            self.MSE_height = tf.reduce_mean(self.error_height ** 2)

            self.error = self.forward[:, :-1] - self.im[:, :-1]  # remaining channels are the actual recon;
            self.MSE = tf.reduce_mean(self.error ** 2)
            self.recon = self.recon[:, :, :-1]  # discard the height map channel, as it's recorded elsewhere;
        else:
            self.error = self.forward - self.im  # save this for computing error map;
            self.MSE = tf.reduce_mean(self.error ** 2)

        if self.report_error_map:
            # project the squared error onto the reconstruction space;
            assert self.momentum is None  # for now, don't use this with batching/momentum;
            # don't have to use interpolated projection, because we don't need to compute gradients for this;

            with tf.device('/CPU:0'):
                rc_warp_int = self.restrict(tf.cast(self.rc_warp, tf.int32))
                self.error_map = tf.scatter_nd(rc_warp_int, self.error ** 2,
                                               [self.recon_shape[0], self.recon_shape[1], self.num_channels])
                norm = tf.scatter_nd(rc_warp_int, tf.ones(self.rc_warp.shape[0], dtype=self.tf_dtype), self.recon_shape)
                self.error_map = tf.math.divide_no_nan(self.error_map, norm[:, :, None])
        else:
            self.error_map = None

    def _add_regularization_loss(self, reg_coefs):

        # tf Variable to be regularized:
        if 'perspective_to_' in self.deformation_model:
            field = self.ego_height_to_regularize

        self.loss_list = [self.MSE]
        if 'L2' in reg_coefs and reg_coefs['L2'] is not None:
            self.loss_list.append(reg_coefs['L2'] * tf.reduce_sum(field ** 2) / self.scale ** 2)
        if 'L1' in reg_coefs and reg_coefs['L1'] is not None:
            self.loss_list.append(reg_coefs['L1'] * tf.reduce_sum(tf.sqrt(field ** 2 + 1e-7)) / self.scale ** 2)
        if 'TV2' in reg_coefs and reg_coefs['TV2'] is not None:  # TV squared;
            d0 = field[:, 1:, :-1] - field[:, :-1, :-1]
            d1 = field[:, :-1, 1:] - field[:, :-1, :-1]
            self.loss_list.append(reg_coefs['TV2'] * tf.reduce_sum(d0 ** 2 + d1 ** 2) / self.scale ** 2)
        if 'TV' in reg_coefs and reg_coefs['TV'] is not None:
            d0 = field[:, 1:, :-1] - field[:, :-1, :-1]
            d1 = field[:, :-1, 1:] - field[:, :-1, :-1]
            self.loss_list.append(reg_coefs['TV'] * tf.reduce_sum(tf.sqrt(d0 ** 2 + d1 ** 2 + 1e-7)) / self.scale ** 2)
        if 'height_map' in reg_coefs and reg_coefs['height_map'] is not None:
            # regularize consistency of deformations to the global height map;
            self.loss_list.append(reg_coefs['height_map'] * self.MSE_height)

    @tf.function
    def gradient_update(self, stack_downsamp, rc_downsamp, update_gradient=True, reg_coefs=None, dither_coords=True,
                        return_tracked_tensors=False, stop_gradient=True, p2p_warp_mode=None, return_grads=False,
                        return_loss_only=False):
        # do one gradient update step; pass thru the downsampled stack and rc coordinates via tf.datasets if batching;
        # one might set update_gradient to False if using batching and just want to create the reconstruction without
        # registering simultaneously;
        # dither_coords: if True, add a random float between -1 and 1 to prevent adaptation to the interpolation scheme;
        # stop_gradient: whether to use tf.stop_gradient on the reconstruction before computing the forward prediction;
        # p2p_warp_mode: see warp_camera_parameters and/or generate_recon; if using perspective-to-orthographic warping,
        # then set this to None; if using perspective-to-perspective warping, then set to one of 'mean', 'random',
        # 'fixed';
        # return_loss_only: as the name suggests;

        with tf.GradientTape() as tape:
            self._generate_recon(stack_downsamp, rc_downsamp, dither_coords, p2p_warp_mode, update_gradient)
            if stop_gradient:
                self.recon = tf.stop_gradient(self.recon)
            self._forward_prediction()

            if reg_coefs is not None:
                self._add_regularization_loss(reg_coefs)
                loss = tf.reduce_sum(self.loss_list)
            else:
                loss = self.MSE

        grads = tape.gradient(loss, self.train_var_list)
        # apply gradient update for each optimizer:
        if update_gradient:
            for grad, var, optimizer, train in zip(grads, self.train_var_list, self.optimizer_list,
                                                   self.trainable_or_not):
                if train:  # if user specifies negative learning rate, then don't update;
                    if type(var) is list or type(var) is ListWrapper:  # sometimes it's a ListWrapper; after restoring
                        # from checkpoint?
                        # this is probably the neural network variable list;
                        optimizer.apply_gradients(zip(grad, var))
                    else:
                        optimizer.apply_gradients([(grad, var)])

            # update other stuff:
            if 'camera_parameters' in self.deformation_model and self.force_ground_surface_up and update_gradient:
                # rotate all vectors such that their mean direction points in [0,0,-1];
                # (mean across all, not just batch!);
                # using rodrigues formula;
                n_mean = tf.reduce_mean(self.ground_surface_normal, axis=0)
                rot_axis = tf.stack([-n_mean[1], n_mean[0], 0])  # cross product;
                rot_axis, _ = tf.linalg.normalize(rot_axis)
                cos_rot_angle = -n_mean[2] / tf.norm(n_mean)
                sin_rot_angle = tf.sqrt(1 - cos_rot_angle ** 2 + 1e-6)  # prevent numerical issues;
                K = tf.stack([[0, -rot_axis[2], rot_axis[1]],
                              [rot_axis[2], 0, -rot_axis[0]],
                              [-rot_axis[1], rot_axis[0], 0]])
                rotmat = tf.eye(3) + sin_rot_angle * K + K @ K * (1 - cos_rot_angle)
                self.ground_surface_normal.assign(tf.einsum('ij,cj->ci', rotmat, self.ground_surface_normal))

        if return_loss_only:
            if reg_coefs is not None:
                return self.loss_list
            else:
                return self.MSE
        else:
            if reg_coefs is not None:
                return_list = [self.loss_list, self.recon, self.normalize, self.error_map]
            else:
                return_list = [self.MSE, self.recon, self.normalize, self.error_map]

            if return_tracked_tensors:
                return_list.append(self.tensors_to_track)
            if return_grads:
                return_list.append(grads)

            return return_list

    def checkpoint_all_variables(self, path='./tf_ckpts', skip_saving=False):
        # save checkpoints to restore in case unet diverges;
        # skip_saving if you just want to create the ckpt and manager;
        if self.ckpt is None:
            self.ckpt = tf.train.Checkpoint()
            self.ckpt.var = self.train_var_list
            self.ckpt.opt = self.optimizer_list
            self.ckpt.non_train_var = self.non_train_list
            self.manager = tf.train.CheckpointManager(self.ckpt, path, max_to_keep=2)
            # only keep two, restore the oldest;
        if not skip_saving:
            self.manager.save()

    def restore_all_variables(self):
        self.ckpt.restore(self.manager.checkpoints[0])

    @tf.function
    def get_all_variables(self):
        # return as dictionary so it can be used as initialization:
        with tf.device('/CPU:0'):
            # don't return neural network parameters -- use checkpoints instead;
            vars = dict()
            for var in self.train_var_list:
                try:
                    vars[var.name[:-2]] = var
                except AttributeError:  # ListWrapper for unet, for which we don't need the variables;
                    pass
            return vars


# u-net for mapping images to height map:
class unet(tf.keras.Model):
    def __init__(self, filters_list, skip_list, output_nonlinearity, upsample_method='bilinear'):
        # filters_list and skip_list are lists of number of filters in the upsample/downsample layers,
        # and the number of filters in the skip connections;
        # output_nonlinearity can be 'leaky_relu' or 'linear';
        # upsample_method can be 'bilinear' or 'nearest', which is used in the upsample blocks;
        super(unet, self).__init__()
        assert len(filters_list) == len(skip_list)
        self.filters_list = filters_list
        self.skip_list = skip_list
        self.output_nonlinearity = output_nonlinearity
        self.upsample_method = upsample_method
        self._build()

    def _build(self):
        # define all the layers of the unet;
        self.downsample_list = list()  # stores list of downsample blocks;
        self.downsample_skip_block_list = list()  # stores list of skip convolutional blocks;
        self.upsample_list = list()  # stores list of upsample blocks;
        self.upsample_concat_list = list()  # stores list of concatenation layers;

        # downsampling half:
        for num_filters, num_skip_filters in zip(self.filters_list, self.skip_list):
            self.downsample_list.append(self._downsample_block(num_filters))  # add to list of layers
            self.downsample_skip_block_list.append(self._skip_block(num_skip_filters))

        # upsampling half:
        for i, (num_filters, num_skip_filters) in enumerate(zip(self.filters_list[::-1], self.skip_list[::-1])):
            if num_skip_filters != 0:
                self.upsample_concat_list.append(tf.keras.layers.Concatenate())
            else:
                self.upsample_concat_list.append(None)  # as a placeholder
            if i == len(self.filters_list) - 1:
                # last block, use the specified output nonlinearity:
                self.upsample_list.append(self._upsample_block(num_filters,
                                                               nonlinearity=self.output_nonlinearity))
            else:
                self.upsample_list.append(self._upsample_block(num_filters))

    def _downsample_block(self, numfilters, kernel_size=3):
        return [tf.keras.layers.Conv2D(filters=numfilters, kernel_size=kernel_size,
                                       strides=(2, 2), padding='same'),
                tf.keras.layers.BatchNormalization(),
                tf.keras.layers.LeakyReLU(),
                # repeat, but no downsample this time
                tf.keras.layers.Conv2D(filters=numfilters, kernel_size=3,
                                       strides=(1, 1), padding='same'),
                tf.keras.layers.BatchNormalization(),
                tf.keras.layers.LeakyReLU()]

    def _upsample_block(self, numfilters, kernel_size=3, nonlinearity='leaky_relu'):
        layers_list = [tf.keras.layers.UpSampling2D(interpolation=self.upsample_method),
                       tf.keras.layers.Conv2D(filters=numfilters, kernel_size=kernel_size,
                                              strides=(1, 1), padding='same'),
                       tf.keras.layers.BatchNormalization(),
                       tf.keras.layers.LeakyReLU(),
                       tf.keras.layers.Conv2D(filters=numfilters, kernel_size=1,
                                              strides=(1, 1), padding='same'),  # kernel size 1
                       tf.keras.layers.BatchNormalization()]
        if nonlinearity == 'leaky_relu':
            layers_list.append(tf.keras.layers.LeakyReLU())
        elif nonlinearity == 'linear':
            pass
        else:
            raise Exception('invalid nonlinearity')
        return layers_list

    def _skip_block(self, numfilters=4, kernel_size=1):
        if numfilters == 0:  # no skip connections
            return None
        else:
            return [tf.keras.layers.Conv2D(filters=numfilters, kernel_size=kernel_size,
                                           strides=(1, 1), padding='same'),
                    tf.keras.layers.BatchNormalization(),
                    tf.keras.layers.LeakyReLU()]

    def call(self, x):
        skip_layers = list()  # store skip layer outputs to be concatenated in the upsample block;
        for down_block, skip_block in zip(self.downsample_list, self.downsample_skip_block_list):
            for down_layer in down_block:  # traverse all layers in block;
                x = down_layer(x)
            if skip_block is not None:  # if there's a skip block, traverse all layers in it;
                x_ = x
                for skip_layer in skip_block:
                    x_ = skip_layer(x_)
                skip_layers.append(x_)
            else:
                skip_layers.append(None)
        for up_block, skip, concat in zip(self.upsample_list, skip_layers[::-1], self.upsample_concat_list):
            if skip is not None:
                x = concat([x, skip])
            for up_layer in up_block:
                x = up_layer(x)
        return x


# this function is used with the unet to calculate how to pad the input (and subsequently depad the output):
def get_compatible_size(dim, num_downsamp, max_dim=10000):
    # for a given dim and number of downsample blocks, find the smallest value >= dim such that the
    # unet will return the same size as the input;
    # max_dim is the largest value to be considered;
    k = 2 ** num_downsamp
    for i in range(1, max_dim):
        new_dim = i * k
        if new_dim >= dim:
            break
    return new_dim

# below are convenience functions for the script to call:
def stack_loader_phone(directory):
    # for my phone datasets;
    im_stack = list()
    for filename in tqdm(sorted(os.listdir(directory))):
        if 'jpg' in filename or 'png' in filename:
            im = plt.imread(directory + filename)
            im_stack.append(im)
    im_stack = np.array(im_stack)
    shape = im_stack.shape[1:]
    print('Shape of stack: ' + str(shape))
    return im_stack

def xcorr_initial_guess(im_stack, downsamp=4, channel=0, crop_frac=.4):
    # generates initial guess for x_s and y_pos using successive cross correlations (therefore, sequential images should
    # have overlap);
    # expect the input im_stack to be the output of an image loading script;
    # specify downsample factor (without antialiasing) and which channel to use;
    # crop_frac -- only search a center crop of the xcorr;

    im_stack_ = im_stack[:, ::downsamp, ::downsamp, channel].astype(np.float32)
    pos = [[0, 0]]  # first coordinate is arbitrarily defined to be the origin;
    for i in range(len(im_stack) - 1):
        xcorr = scipy.signal.correlate(im_stack_[i], im_stack_[i + 1])
        bkgd = scipy.signal.correlate(im_stack_[i], im_stack_[i])  # use autocorrelation as background;
        xcorr -= bkgd  # subtract background to avoid bias approaching 0 delay
        s = np.array(xcorr.shape)
        start_pix = np.int32((s - s * crop_frac) / 2)
        end_pix = np.int32((s + s * crop_frac) / 2)
        xcorr = xcorr[start_pix[0]:end_pix[0], start_pix[1]:end_pix[1]]

        argmax = np.unravel_index(xcorr.argmax(), xcorr.shape)
        pos.append(argmax - np.array(xcorr.shape) / 2)
    pos = np.array(pos).astype(np.float32) * downsamp
    pos = np.cumsum(pos, axis=0)  # convert relative shifts into absolute positions
    x_pos = pos[:, 1]
    y_pos = pos[:, 0]
    x_pos -= np.min(x_pos)  # avoid negative values because they will be stored as uints
    y_pos -= np.min(y_pos)

    return x_pos, y_pos

def monitor_progress(recon, error_map, losses, tracked=None):
    # run this in your train loop to print out progress;
    # recon is the reconstruction, error_map is the error map, losses is a list of MSEs or loss terms, and tracked
    # is a dictionary of tracked tensors from the tf graph;

    num_channels = recon.shape[-1]
    plt.figure(figsize=(15, 15))
    if num_channels == 1:
        plt.imshow(recon[:, :, 0].numpy(), cmap='jet')
        plt.colorbar()
    else:
        plt.imshow(np.uint8(np.clip(recon.numpy(), 0, 255)))
    plt.title('Reconstruction')
    plt.show()

    if error_map is not None:
        plt.figure(figsize=(15, 15))
        plt.imshow(error_map.numpy().sum(2))
        plt.clim([0, 300])
        plt.colorbar()
        plt.title('Error map')
        plt.show()

    if tracked is not None and 'height_map' in tracked:
        height_map = tracked['height_map']
        mask = recon.numpy().sum(2) > 0  # which pixels were visited (also could use normalize);
        clims = np.percentile(height_map[mask], [.1, 99.9])
        plt.figure(figsize=(15, 15))
        plt.imshow(height_map, cmap='jet')
        plt.clim(clims)
        plt.title('Height map')
        plt.colorbar()
        plt.show()

    plt.plot(losses)
    plt.title('Loss history')
    plt.show()

import numpy as np
import os
import matplotlib.pyplot as plt
import cv2
import tensorflow as tf
from tqdm.notebook import tqdm
import scipy.signal
from tensorflow.python.training.tracking.data_structures import ListWrapper


class mesoSfM:
    def __init__(self, stack, ul_coords, recon_shape, ul_offset=(0, 0), batch_size=None, scale=1, momentum=None,
                 batch_across_images=False, report_error_map=False, restrict_function='clip'):
        # stack: stack of images to be stitched; 3D tensor of shape num_images, num_rows, num_cols, num_channels;
        # ul_coords: array of upper-left coordinates of the images (in pixels); 2D tensor of shape num_images, 2 (row,
        # col; y, x);
        # recon_shape: shape of the final reconstruction/stitched image (row, col, channels);
        # ul_offset: length-2 vector (row, col) specifying a constant offset for all the ul_coords (in pixels);
        # batch_size: if None, then don't use batching;
        # scale: factor between 0 and 1 specifying how much to downsample the reconstruction;
        # momentum: only relevant for batching; specifies how much averaging of previous iterations to use for recon;
        # either both batch_size and momentum must be non-None, or both must be None;
        # batch_across_images: if you're batching (i.e., batch_size is not None), then if true, then batch across the
        # image dimension; otherwise, batch across pixels; (if batching while using unet, then you must batch across
        # image dimension);
        # report_error_map:
        # restrict_function choice decides what to do if a point goes beyond the boundaries; 'clip' or 'mod' or
        # 'mod_with_random_shifts';

        self.tf_dtype = tf.float32
        self.np_dtype = np.float32

        self.stack = np.uint8(stack)  # cast is 8-bit to save memory for the batch generator; cast batch to float32;
        self.num_channels = self.stack.shape[3]  # number of channels; stack must at least have a singleton dim 3;
        self.num_images = self.stack.shape[0]  # number of images in dataset;
        self.ul_coords = np.uint16(ul_coords)  # uint8 is too narrow, because these pixel coordinates are large;
        self.recon_shape_base = recon_shape  # this will be the base recon_shape; the effective recon_shape will depend
        # on the scale factor that the user specifies;
        self.ul_offset = np.array(ul_offset)
        self.scale = scale
        self.momentum = momentum
        self.batch_size = batch_size
        self.batch_across_images = batch_across_images
        self.report_error_map = report_error_map
        self.restrict_function = restrict_function
        self.sig_proj = .42465  # for the intepolation kernel width;
        self.subtract_min_from_height_map = True
        self.optimizer = tf.keras.optimizers.Adam

        # unet parameters if relevant; define these manually if needed;
        self.filters_list = None
        self.skip_list = None
        self.unet_scale = .01  # to scale the output of the unet
        self.output_nonlinearity = 'linear'  # 'linear' or 'leaky_relu'
        self.upsample_method = 'bilinear'  # 'bilinear' or 'nearest'
        self.ckpt = None  # for checkpointing models when using unet;
        self.save_iter = 15  # save model every this many iterations;
        self.recompute_CNN = False  # save memory using tf.recompute_grad;

        self.height_scale_factor = 4000  # to scale the ego_height parameter to avoid dealing with large values;

        # camera parameters for getting absolute scale;
        self.use_absolute_scale_calibration = False
        self.effective_focal_length_mm = 4.3  # effective focal length in mm;
        self.magnification_j = None  # user needs to figure out the magnification of the jth camera view;
        self.j = 0  # and specify which camera view index; by default, assume the first camera;

    def create_variables(self, deformation_model, learning_rates=None, variable_initial_values=None, recon=None,
                         normalize=None, remove_global_transform=False, antialiasing_filter=False,
                         stack_downsample_factor=None, force_ground_surface_up=False):
        # define tf.Variables and optimizers;
        # deformation_model: affects what tf.Variables will be defined and optimized;
        # learning_rates is a dictionary of variable names (strings) and the corresponding learning rate; None means
        # use the default, defined below; can specify a learning rate a negative number, indicating that variable is not
        # optimized; can supply initial variable values as a dictionary; if None, use default;
        # if recon and normalize are supplied, then initialize recon_previous and normalize_previous with these after
        # upsampling; both must be supplied; (this is for momentum/running average calculations); most likely, recon and
        # normalize will be from earlier reconstruction attempts using mesoSfM, but at lower resolution;
        # remove_global_transform: e.g., no global shift or scale;
        # stack_downsample_factor is an integer that downsamples the stack and coordinates; if None, then it will be
        # computed from self.scale;
        # force_ground_surface_up; when using camera model, force the mean surface normal to be [0,0,-1];

        # these must be both be None, or neither None:
        assert (self.batch_size is not None) == (self.momentum is not None)
        assert (recon is None) == (normalize is None)

        # define downsample factor:
        self.recon_shape = np.int32(self.recon_shape_base * self.scale)
        if stack_downsample_factor is None:
            self.downsample = np.int32(1 / self.scale)  # also downsample the images to save computation;
            self.downsample = np.maximum(self.downsample, 1)  # obviously can't downsample with 0;
        else:
            self.downsample = stack_downsample_factor
        self.im_downsampled_shape = np.array([(self.stack.shape[1] - 1) // self.downsample + 1,
                                              (self.stack.shape[2] - 1) // self.downsample + 1])

        # coordinates of the images:
        c = np.arange(self.stack.shape[2], dtype=np.uint16)
        r = np.arange(self.stack.shape[1], dtype=np.uint16)
        r, c = np.meshgrid(r, c, indexing='ij')
        rc_base = np.stack([r, c]).T
        self.rc_base = np.tile(rc_base[None],
                               [self.num_images, 1, 1, 1])  # base coordinates (could be downsampled after applying a
        # scale factor);

        self.deformation_model = deformation_model
        self.remove_global_transform = remove_global_transform
        self.force_ground_surface_up = force_ground_surface_up  # only relevant if using a camera model;

        # list of tf.Variables and optimizers, to be populated by one or more _create_variables_*deformation_model*;
        self.train_var_list = list()
        self.optimizer_list = list()
        self.non_train_list = list()  # list of variables that aren't trained (probably .assign()'d; for checkpoints);
        self.tensors_to_track = dict()  # intermediate tensors to track; have a tf.function return the contents;

        def use_default_for_missing(input_dict, default_dict):
            # to be used directly below; allows for dictionaries in which not all keys are specified; if not specified,
            # then use default_dict's value;

            if input_dict is None:  # if nothing given, then use the default;
                return default_dict
            else:
                for key in default_dict:
                    if key in input_dict:
                        if input_dict[key] is None:  # if the key is present, but None is specified;
                            input_dict[key] = default_dict[key]
                        else:  # i.e., use the value given;
                            pass
                    else:  # if key is not even present;
                        input_dict[key] = default_dict[key]
                return input_dict

        if 'camera_parameters' in deformation_model:
            if 'unet' in deformation_model:
                # user has to manually define these:
                assert self.filters_list is not None
                assert self.skip_list is not None

            default_learning_rates = {'camera_focal_length': 1e-3, 'camera_height': 1e-3, 'ground_surface_normal': 1e-3,
                                      'camera_in_plane_angle': 1e-3, 'rc': 10, 'gain': 1e-3, 'ego_height': 1e-3,
                                      'bias': 1e-3}
            default_variable_initial_values = {'camera_focal_length': np.float32(1),  # initialize same value as ...
                                               'camera_height': 1 * np.ones([self.num_images]),  # ...for height;
                                               'ground_surface_normal': np.concatenate(
                                                   [np.zeros((self.num_images, 2)) + 1e-7,  # to avoid /0;
                                                    -np.ones((self.num_images, 1))], axis=1),
                                               'camera_in_plane_angle': np.zeros([self.num_images]),
                                               'rc': self.ul_coords,
                                               'ego_height': 1e-7 +  # add small value to allow gradients to prop;
                                                         np.zeros([self.num_images,
                                                                   self.im_downsampled_shape[0],
                                                                   self.im_downsampled_shape[1]]),
                                               'gain': np.ones(self.num_images),
                                               'bias': np.zeros(self.num_images)}
            learning_rates = use_default_for_missing(learning_rates, default_learning_rates)
            variable_initial_values = use_default_for_missing(variable_initial_values, default_variable_initial_values)
            self._create_variables_camera_parameters(learning_rates, variable_initial_values)

            # these are for perspective to perspective and perspective to orthographic:
            if 'unet' not in deformation_model:
                # create radial deformation field:
                self._create_variables_perspective_to_orthographic(learning_rates, variable_initial_values)
            elif 'unet' in deformation_model:
                # create neural network layers:
                self._create_variables_height_map_unet(learning_rates, variable_initial_values)

            if self.remove_global_transform:
                # removing the global scaling transform entails adapting the focal length to the mean height;
                assert learning_rates['camera_focal_length'] < 0

            if 'perspective_to_perspective' in deformation_model:
                # these won't be optimized; also these initializing values are very temporary, and should be immediately
                # modified by code later; the reason why they are tf.Variables is so that the user can manually modify
                # them in eager mode with .assign();
                self.reference_camera_height = tf.Variable(1, dtype=self.tf_dtype, name='reference_camera_height')
                self.reference_camera_rc = tf.Variable(np.zeros(2), dtype=self.tf_dtype, name='reference_camera_rc')
                # note that because there's no tilt, reference_camera_height is the same as the camera to vanishing
                # point distance, and the reference_camera_rc is the same as the vanishing point position;
                self.non_train_list.append(self.reference_camera_rc)
                self.non_train_list.append(self.reference_camera_height)
        else:
            raise Exception('invalid deformation model: ' + deformation_model)

        # intensity adjustment (uniform within a given image for now):
        self.gain = tf.Variable(variable_initial_values['gain'], dtype=self.tf_dtype, name='gain')
        self.gain_optim = self.optimizer(learning_rate=learning_rates['gain'])
        self.train_var_list.append(self.gain)
        self.optimizer_list.append(self.gain_optim)

        # intensity bias (also uniform within each image):
        self.bias = tf.Variable(variable_initial_values['bias'], dtype=self.tf_dtype, name='bias')
        self.bias_optim = self.optimizer(learning_rate=learning_rates['bias'])
        self.train_var_list.append(self.bias)
        self.optimizer_list.append(self.bias_optim)

        # barrel or pincushion distortion correction:
        if 'radial_camera_distortion' in learning_rates:
            assert 'camera' in deformation_model  # this is explicitly a camera modeling option;
            if 'radial_camera_distortion' in variable_initial_values:
                camera_distortion_init = variable_initial_values['radial_camera_distortion']
                if np.ndim(camera_distortion_init) == 0:  # user is specifying there to be only one common parameter;
                    camera_distortion_init = np.reshape(camera_distortion_init, (1, 1))  # make sure at least 2 dims;
                elif np.ndim(camera_distortion_init) == 1:
                    # assume this one dimension refers to the camera view, not the polynomial order;
                    camera_distortion_init = np.reshape(camera_distortion_init, (-1, 1))
            else:
                camera_distortion_init = np.zeros((self.num_images, 1))  # second dim: polynomial order;
            num_poly_terms = camera_distortion_init.shape[1]  # number of terms in the polynomial (even only);
            self.radial_powers = (np.arange(num_poly_terms) + 1)  # half of the even powers to raise to;
            self.correct_radial_camera_distortion = camera_distortion_init.shape[0]  # this info used if batching;
            self.radial_camera_distortion = tf.Variable(camera_distortion_init, dtype=self.tf_dtype,
                                                        name='radial_camera_distortion')
            self.radial_camera_distortion_optim = self.optimizer(
                learning_rate=learning_rates['radial_camera_distortion'])
            self.train_var_list.append(self.radial_camera_distortion)
            self.optimizer_list.append(self.radial_camera_distortion_optim)
        else:
            self.correct_radial_camera_distortion = False

        if 'radial_camera_distortion_piecewise_linear' in learning_rates:
            assert 'camera' in deformation_model
            if 'radial_camera_distortion_piecewise_linear' in variable_initial_values:
                camera_distortion_init = variable_initial_values['radial_camera_distortion_piecewise_linear']
                assert np.ndim(camera_distortion_init) == 1  # for now, only allow a common distortion among all cams;
                # length of this vector determines how many discretization levels;
            else:
                camera_distortion_init = np.zeros(50)
            self.num_radial_pixels = len(camera_distortion_init)  # how many discretization levels (nodes);
            self.correct_radial_camera_distortion_piecewise_linear = -1  # used if batching images; this never equals
            # num_images;
            self.radial_camera_distortion_piecewise_linear = tf.Variable(camera_distortion_init, dtype=self.tf_dtype,
                                                                         name='radial_camera_distortion_piecewise_linear')
            self.radial_camera_distortion_piecewise_linear_optim = self.optimizer(
                learning_rate=learning_rates['radial_camera_distortion_piecewise_linear'])
            self.train_var_list.append(self.radial_camera_distortion_piecewise_linear)
            self.optimizer_list.append(self.radial_camera_distortion_piecewise_linear_optim)
        else:
            self.correct_radial_camera_distortion_piecewise_linear = False

        # if camera is not centered; this defines the center of the above two distortions, so center the camera first,
        # apply the distortions, then decenter back;
        if 'camera_distortion_center' in learning_rates:
            assert 'camera' in deformation_model
            if 'camera_distortion_center' in variable_initial_values:
                camera_distortion_center_init = variable_initial_values['camera_distortion_center']
                if np.ndim(camera_distortion_center_init) == 1:  # one common pair of parameters;
                    assert len(camera_distortion_center_init) == 2  # x and y centers;
                    camera_distortion_center_init = tf.reshape(camera_distortion_center_init, (1, -1))
                if np.ndim(camera_distortion_center_init) == 0:
                    raise Exception('must supply two values for x/y centration parameters')
            else:
                camera_distortion_center_init = np.zeros((self.num_images, 2))  # by default, diff pair for each camera;
            self.correct_camera_distortion_center = camera_distortion_center_init.shape[0]  # this info used if
            # batching;
            self.camera_distortion_center = tf.Variable(camera_distortion_center_init, dtype=self.tf_dtype,
                                                    name='camera_distortion_center')
            self.camera_distortion_center_optim = self.optimizer(
                learning_rate=learning_rates['camera_distortion_center'])
            self.train_var_list.append(self.camera_distortion_center)
            self.optimizer_list.append(self.camera_distortion_center_optim)
        else:
            self.correct_camera_distortion_center = False

        # create a list of booleans to accompany self.train_var_list and self.optimizer_list to specify whether to train
        # those variables (as specified by the whether the user-specified learning rates are negative); doing this so
        # that autograph doesn't traverse all branches of the conditionals; if the user ever wants to turn off
        # optimization of a variable mid-optimization, then just do .assign(0) to the learning rate, such that the
        # update is still happening, but the change is 0;
        self.trainable_or_not = list()
        for var in self.train_var_list:
            if type(var) is list:
                # if the variable is a list of variables, then this should be for the unet; modify here if there are
                # other scenarios;
                assert 'unet' in self.deformation_model
                name = 'ego_height'
            else:
                name = var.name[:-2]
            flag = learning_rates[name] > 0
            self.trainable_or_not.append(flag)

        # downsample rc coordinates and stack:
        rc = np.transpose(self.rc_base[:, ::self.downsample, ::self.downsample, :], (0, 2, 1, 3))
        if antialiasing_filter:
            downsample = int(self.downsample)  # cv2 doesn't like numpy values?
            if downsample == 1:
                print('warning: antialiasing filter is applied even though there is no downsampling')
            ksize = int(downsample * 2.5)
            if ksize % 2 == 0:  # must be odd;
                ksize += 1
            stack_downsamp = np.stack(
                [cv2.GaussianBlur(im, (ksize, ksize), downsample, downsample)
                 [::downsample, ::downsample] for im in self.stack])
        else:
            stack_downsamp = self.stack[:, ::self.downsample, ::self.downsample, :]

        rc_downsamp = np.reshape(rc, (self.rc_base.shape[0], -1, self.rc_base.shape[-1]))  # flatten spatial dims;
        stack_downsamp = np.reshape(stack_downsamp,
                                    [self.num_images, -1, self.num_channels])  # flatten spatial dims;
        self.rc_downsamp = rc_downsamp
        self.stack_downsamp = stack_downsamp

        # create variables relevant for batching (or set them to None if not batching):
        if self.momentum is not None:
            if 'camera_parameters_perspective_' in deformation_model:
                # we're going to give the coregistered height map a ride;
                num_channels = self.num_channels + 1
            else:
                num_channels = self.num_channels

            if recon is None:  # if none supplied, initialize with 0s;
                recon_previous = np.zeros([self.recon_shape[0], self.recon_shape[1], num_channels])
                normalize_previous = np.zeros(self.recon_shape)
            else:  # otherwise, upsample to the current shape;
                recon_previous = cv2.resize(np.nan_to_num(recon), tuple(self.recon_shape[::-1]))
                normalize_previous = cv2.resize(np.nan_to_num(normalize), tuple(self.recon_shape[::-1]))
                if num_channels == 1:
                    # cv2 seems to squeeze singleton channels dimensions, so if it's singleton, add it back:
                    recon_previous = recon_previous[:, :, None]
                if recon_previous.shape[-1] != num_channels:
                    # if you would like to use the previous RGB image as initialization, but your previous run didn't
                    # estimate a height map;
                    assert 'camera_parameters_perspective_' in deformation_model
                    # add empty height channel:
                    recon_previous = np.concatenate([recon_previous, np.zeros_like(recon_previous[:, :, 0:1])], axis=-1)

            # initialize first recon and normalize tensors for momentum; use the scaled recon shape, not the base shape;
            self.recon_previous = tf.Variable(recon_previous, dtype=self.tf_dtype, trainable=False)
            self.non_train_list.append(self.recon_previous)
        else:
            self.recon_previous = None
            self.normalize_previous = None

    def _create_variables_perspective_to_orthographic(self, learning_rates, variable_initial_values):
        # radially inwardly pointing vector magnitudes, where the larger the magnitude, the taller the object;

        ego_height = variable_initial_values['ego_height']
        # make sure the first dimensions match (number of images in stack):
        assert ego_height.shape[0] == self.num_images
        if ego_height.shape[1:] != tuple(self.im_downsampled_shape):
            # presumably you've initialized this with the results from another optimization at a different scale;
            # thus, resize to match the current scale:
            if type(ego_height) != np.ndarray:
                # convert from tf to np if you need to:
                ego_height = ego_height.numpy()
            ego_height = np.stack([cv2.resize(im, tuple(self.im_downsampled_shape[::-1])) for im in ego_height])

        self.ego_height = tf.Variable(ego_height, dtype=self.tf_dtype, name='ego_height')
        self.ego_height_optim = self.optimizer(learning_rate=learning_rates['ego_height'])
        self.train_var_list.append(self.ego_height)
        self.optimizer_list.append(self.ego_height_optim)

    def _create_variables_height_map_unet(self, learning_rates, variable_initial_values):

        self.network = unet(self.filters_list, self.skip_list, output_nonlinearity=self.output_nonlinearity,
                            upsample_method=self.upsample_method)
        # run the network once so that we can access network.trainable_variables
        self.network(tf.zeros([1, 2 ** len(self.filters_list), 2 ** len(self.filters_list), self.num_channels],
                              dtype=self.tf_dtype))

        self.train_var_list.append(self.network.trainable_variables)
        self.optimizer_list.append(self.optimizer(learning_rate=learning_rates['ego_height']))

        if self.recompute_CNN:
            self.network = tf.recompute_grad(self.network)

        # get padded shape that the network likes;
        self.padded_shape = [get_compatible_size(dim, len(self.filters_list)) for dim in self.im_downsampled_shape]
        pad_r = self.padded_shape[0] - self.im_downsampled_shape[0]
        pad_c = self.padded_shape[1] - self.im_downsampled_shape[1]
        pad_top = pad_r // 2
        pad_bottom = int(tf.math.ceil(pad_r / 2))
        pad_left = pad_c // 2
        pad_right = int(tf.math.ceil(pad_c / 2))
        pad_specs = ((pad_top, pad_bottom), (pad_left, pad_right))
        self.pad_layer = tf.keras.layers.ZeroPadding2D(pad_specs)
        self.depad_layer = tf.keras.layers.Cropping2D(pad_specs)

    def _create_variables_camera_parameters(self, learning_rates, variable_initial_values):
        # coordinates are in camera space;

        self.camera_focal_length = tf.Variable(variable_initial_values['camera_focal_length'], dtype=self.tf_dtype,
                                               name='camera_focal_length')
        self.camera_focal_length_optim = self.optimizer(learning_rate=learning_rates['camera_focal_length'])

        self.camera_height = tf.Variable(variable_initial_values['camera_height'], dtype=self.tf_dtype,
                                         name='camera_height')  # height from the camera perspective;
        self.camera_height_optim = self.optimizer(learning_rate=learning_rates['camera_height'])

        self.ground_surface_normal = tf.Variable(variable_initial_values['ground_surface_normal'], dtype=self.tf_dtype,
                                                 name='ground_surface_normal')
        self.ground_surface_normal_optim = self.optimizer(
            learning_rate=learning_rates['ground_surface_normal'])

        self.camera_in_plane_angle = tf.Variable(variable_initial_values['camera_in_plane_angle'],
                                                 dtype=self.tf_dtype, name='camera_in_plane_angle')
        self.camera_in_plane_angle_optim = self.optimizer(
            learning_rate=learning_rates['camera_in_plane_angle'])

        self.rc_ul_per_im = tf.Variable(variable_initial_values['rc'], dtype=self.tf_dtype, name='rc')
        self.rc_ul_per_im_optim = self.optimizer(learning_rate=learning_rates['rc'])

        self.train_var_list.append(self.camera_focal_length)
        self.optimizer_list.append(self.camera_focal_length_optim)
        self.train_var_list.append(self.camera_height)
        self.optimizer_list.append(self.camera_height_optim)
        self.train_var_list.append(self.ground_surface_normal)
        self.optimizer_list.append(self.ground_surface_normal_optim)
        self.train_var_list.append(self.camera_in_plane_angle)
        self.optimizer_list.append(self.camera_in_plane_angle_optim)
        self.train_var_list.append(self.rc_ul_per_im)
        self.optimizer_list.append(self.rc_ul_per_im_optim)

    def generate_dataset(self):
        # user calls this function to get a dataset to iterate over; if not using batching, then just return a tuple or
        # list of length 1 (i.e., the whole dataset is one batch);

        if self.batch_size is not None:
            if self.batch_across_images:
                # sample a subset of the images, and keep track of the indices downsampled so that you can gather the
                # corresponding variables;
                tensor_slices = (self.stack_downsamp,
                                 (self.rc_downsamp, np.arange(self.num_images, dtype=np.int32)))
                dataset = (tf.data.Dataset.from_tensor_slices(tensor_slices).shuffle(self.num_images)
                           .batch(self.batch_size, drop_remainder=True).repeat(None).prefetch(1))
                return dataset
            else:
                # transpose to batch along space, not image number;
                rc_downsamp_T = np.transpose(self.rc_downsamp, (1, 0, 2))
                stack_downsamp_T = np.transpose(self.stack_downsamp, (1, 0, 2))
                if 'camera_parameters_perspective_' in self.deformation_model:
                    # need to also get coordinates of the spatial positions to index into pixel-wise deformation fields:
                    tensor_slices = (stack_downsamp_T,
                                     (rc_downsamp_T, np.arange(np.prod(self.im_downsampled_shape), dtype=np.int32)))
                else:
                    tensor_slices = (stack_downsamp_T, rc_downsamp_T)
                dataset = (tf.data.Dataset.from_tensor_slices(tensor_slices)
                           .shuffle(len(rc_downsamp_T)).batch(self.batch_size).repeat(None).prefetch(1))
                return dataset
        else:  # basically a 1-batch dataset;
            return self.stack_downsamp, self.rc_downsamp

    def _warp_camera_parameters(self, rc_downsamp, use_radial_deformation, p2p_warp_mode=None,
                                inds_downsamp=None, stack_downsamp=None):
        # shape of rc_downsamp: num_images, _, 2;
        # use_radial_deformation is a boolean flag specifying whether to do the per-pixel radial deformation fields to
        # warp perspective to othographic OR perspective to perspective; if the latter, then p2p_warp_mode specifies
        # how to specify the perspective reference to warp to; the options are 'mean', 'random', 'fixed', and None,
        # where None means you're using perspective-to-orthographpic warping and not perspective to perspective;
        # inds_downsamp is passed if batching and using radial deformations;
        # stack_downsamp is only needed if using a unet;

        if p2p_warp_mode is None:
            assert 'perspective_to_orthographic' in self.deformation_model
        else:
            assert 'perspective_to_perspective' in self.deformation_model

        if self.batch_across_images and self.batch_size is not None:
            # in generate_recon, we defined the batch and non-batch versions;

            rc_ul_per_im = self.rc_ul_per_im_batch
            gain = self.gain_batch
            bias = self.bias_batch
            camera_height = self.camera_height_batch
            ground_surface_normal = self.ground_surface_normal_batch
            camera_in_plane_angle = self.camera_in_plane_angle_batch
            if 'unet' not in self.deformation_model:
                ego_height = self.ego_height_batch
                self.ego_height_to_regularize = self.ego_height_batch
            else:
                pass
            if self.correct_radial_camera_distortion:
                if self.correct_radial_camera_distortion == self.num_images:
                    radial_camera_distortion = self.radial_camera_distortion_batch
                else:
                    radial_camera_distortion = self.radial_camera_distortion
            if self.correct_radial_camera_distortion_piecewise_linear:
                if self.correct_radial_camera_distortion_piecewise_linear == self.num_images:
                    radial_camera_distortion_piecewise_linear = self.radial_camera_distortion_piecewise_linear_batch
                else:
                    radial_camera_distortion_piecewise_linear = self.radial_camera_distortion_piecewise_linear
            if self.correct_camera_distortion_center:
                if self.correct_camera_distortion_center == self.num_images:
                    camera_distortion_center = self.camera_distortion_center_batch
                else:
                    camera_distortion_center = self.camera_distortion_center
            num_images = self.batch_size  # for reshaping below;
            camera_focal_length = self.camera_focal_length  # not batched;
        else:
            rc_ul_per_im = self.rc_ul_per_im
            gain = self.gain
            bias = self.bias
            camera_height = self.camera_height
            ground_surface_normal = self.ground_surface_normal
            camera_in_plane_angle = self.camera_in_plane_angle
            if 'unet' not in self.deformation_model:
                ego_height = self.ego_height
                self.ego_height_to_regularize = self.ego_height
            if self.correct_radial_camera_distortion:
                radial_camera_distortion = self.radial_camera_distortion
            if self.correct_radial_camera_distortion_piecewise_linear:
                radial_camera_distortion_piecewise_linear = self.radial_camera_distortion_piecewise_linear
            if self.correct_camera_distortion_center:
                camera_distortion_center = self.camera_distortion_center

            num_images = self.num_images  # for reshaping below;
            camera_focal_length = self.camera_focal_length

        if self.remove_global_transform:
            # don't use self.camera_focal_length; set to geometric mean of the camera heights;
            # also, use the tf.Variable version always (not the batch version);
            camera_focal_length = tf.reduce_prod(self.camera_height, axis=0, keepdims=False) ** (1 / self.num_images)
            self.camera_focal_length.assign(camera_focal_length)
            camera_in_plane_angle = camera_in_plane_angle - tf.reduce_mean(self.camera_in_plane_angle)
            rc_ul_per_im = rc_ul_per_im - tf.reduce_mean(self.rc_ul_per_im, axis=0, keepdims=True)

        im_dims = np.array(self.stack.shape)[1:3]  # for normalization of image coordinates;
        max_dim = np.max(im_dims)  # to keep isotropic;
        camera_yx = (rc_downsamp - .5 * im_dims[None, None, :]) / max_dim  # normalize to -.5 to .5;

        if self.correct_camera_distortion_center:
            camera_yx -= camera_distortion_center[:, None, :]

        if self.correct_radial_camera_distortion:
            camera_r2 = camera_yx[:, :, 0] ** 2 + camera_yx[:, :, 1] ** 2  # radial distance squared;
            # dims^: camera, pixels
            camera_r2 *= 2  # make it go from -1 to 1 rather than -.5 to .5;

            if self.correct_radial_camera_distortion:
                # even polynomial to account for distortion:
                camera_even_poly = tf.math.pow(camera_r2[:, :, None], self.radial_powers[None, None, :])
                # dims^: camera, pixels, power
                camera_even_poly = tf.reduce_sum(camera_even_poly * radial_camera_distortion[:, None, :], 2)
                # dims^: camera, pixels
                radial_correction_factor = 1 + camera_even_poly[:, :, None]
                self.tensors_to_track['camera_distortion_radial'] = radial_correction_factor
            else:
                radial_correction_factor = 1
            camera_yx = camera_yx * radial_correction_factor

        if self.correct_radial_camera_distortion_piecewise_linear:
            camera_r = tf.sqrt(camera_yx[:, :, 0] ** 2 + camera_yx[:, :, 1] ** 2)  # radial distance;
            # dims^: camera, pixels; these radial coordinates should be between 0 and .5*sqrt(2), but could go higher
            # if th center moves); thus to be safe, just multiply by num_radial_pixels;
            r_scale = camera_r * self.num_radial_pixels
            # find nearest pixels and distances thereto:
            r_floor = tf.floor(r_scale)
            r_ceil = tf.minimum(r_floor + 1, self.num_radial_pixels - 1)  # to prevent out of range indexing;
            r_middle = r_scale - r_floor
            r_floor = tf.cast(r_floor, dtype=tf.int32)
            r_ceil = tf.cast(r_ceil, dtype=tf.int32)
            distortion = 1 + radial_camera_distortion_piecewise_linear
            distortion /= tf.reduce_max(distortion)  # to prevent global expansion;
            correction_factor_floor = tf.gather(distortion, r_floor)
            correction_factor_ceil = tf.gather(distortion, r_ceil)
            # bilinear interpolation:
            correction_factor = correction_factor_ceil * r_middle + correction_factor_floor * (1 - r_middle)
            camera_yx *= correction_factor[:, :, None]
            self.tensors_to_track['camera_distortion_piecewise_linear'] = correction_factor

        if self.correct_camera_distortion_center:
            camera_yx += camera_distortion_center[:, None, :]

        # in-plane rotation:
        cos = tf.cos(camera_in_plane_angle)
        sin = tf.sin(camera_in_plane_angle)
        rotmat_xy = tf.stack([[cos, sin], [-sin, cos]])
        camera_yx = tf.einsum('cri,ijc->crj', camera_yx, rotmat_xy)

        n_ground, _ = tf.linalg.normalize(ground_surface_normal, axis=1)  # normalize to unit mag;
        # shape^: num_images, 3

        # projecting to object space (computed analytically and taylor-expanded);
        nx = n_ground[:, 0][:, None, None]
        ny = n_ground[:, 1][:, None, None]
        nx2 = nx ** 2
        ny2 = ny ** 2
        x = camera_yx[:, :, 1][:, :, None]
        y = camera_yx[:, :, 0][:, :, None]
        h = camera_height[:, None, None]
        f = camera_focal_length
        # using a taylor expansion:
        flat_x = (h * x / f +
                  h * x * (nx * x + ny * y) / f ** 2 +
                  h * (f ** 2 * nx2 * x + 2 * nx2 * x ** 3 + f ** 2 * nx * ny * y + 4 * nx * ny * x ** 2 * y +
                       2 * ny2 * x * y ** 2) / 2 / f ** 3)
        flat_y = (h * y / f +
                  h * y * (nx * x + ny * y) / f ** 2 +
                  h * (f ** 2 * ny2 * y + 2 * ny2 * y ** 3 + f ** 2 * nx * ny * x + 4 * nx * ny * y ** 2 * x +
                       2 * nx2 * y * x ** 2) / 2 / f ** 3)
        flat_xy = tf.concat([flat_x, flat_y], axis=2)

        n_dot_r = n_ground[:, 2] * camera_height  # shape: num_images; dot product between a point on the
        # ground, r (0, 0, camera_height); (this is needed below);
        self.flat_xy = flat_xy

        if use_radial_deformation:
            # compute the vanishing point, which will be used if you use the projective to orthographic mode;
            vanish_xyz = n_dot_r[:, None] * n_ground  # shape: num_images, 3;
            camera_to_vanish_point_xyz = tf.norm(vanish_xyz, axis=1)  # distance from camera to the ground; the actual
            # height of the camera;

            # projection to object space simplifies to this:
            vanish_xy = -camera_height[:, None] * n_ground[:, :2]

            # vanishing point in camera plane:
            vanish_camera_xyz = n_ground * camera_focal_length / n_ground[:, 2:]  # follow surface normal to camera
            # plane;
            vanish_camera_xy = vanish_camera_xyz[:, 0:2]  # don't need z, it's just the focal length;
            # account for in-plane camera rotation:
            vanish_camera_xy = tf.einsum('ci,ijc->cj', vanish_camera_xy, rotmat_xy)
            self.vanish_xy = vanish_xy

        # convert back to row-column:
        flat_rc = (flat_xy[:, :, ::-1] * max_dim + .5 * im_dims[None, None, :])
        if use_radial_deformation:
            vanish_rc = (vanish_xy[:, ::-1] * max_dim + .5 * im_dims[None, :])
            vanish_camera_rc = (vanish_camera_xy[:, ::-1] * max_dim + .5 * im_dims[None, :])
            self.camera_to_vanish_point_rc = camera_to_vanish_point_xyz * max_dim  # convert from xy units to rc units;
            self.tensors_to_track['camera_to_vanish_point_xyz'] = self.camera_to_vanish_point_rc

        # add translations (same as for the homography and affine implementations):
        # these translations don't affect camera_to_vanish_point;
        rc_warp = rc_ul_per_im[:, None, :] + flat_rc
        if use_radial_deformation:
            vanish_warp = rc_ul_per_im + vanish_rc
        if self.restrict_function == 'mod_with_random_shifts':
            # to discourage registration with overlapped regions;
            random_shift = tf.random.uniform(shape=(1, 1, 2), minval=0, maxval=self.recon_shape.max() / self.scale)
            rc_warp += random_shift
            if use_radial_deformation:
                vanish_warp += random_shift[0]
        else:
            rc_warp += self.ul_offset[None, None, :]
            if use_radial_deformation:
                vanish_warp += self.ul_offset[None, :]

        if 'unet' in self.deformation_model:
            # generate self.ego_height; doesn't matter if batching or not because it's generated from the image batch;
            unet_input = tf.reshape(stack_downsamp, [num_images, self.im_downsampled_shape[0],
                                                     self.im_downsampled_shape[1], self.num_channels])

            unet_input = self.pad_layer(unet_input)
            unet_output = self.network(unet_input)
            ego_height = tf.reduce_mean(self.depad_layer(unet_output), [-1])  # remove last dimension;
            ego_height *= self.unet_scale
            self.ego_height_to_regularize = ego_height  # need this for regularization;
            self.ego_height = ego_height

        # multiplicative version:
        if use_radial_deformation:
            self.tensors_to_track['vanish_warp'] = vanish_warp
            self.tensors_to_track['vanish_camera'] = vanish_camera_rc
            if p2p_warp_mode is None:  # using perspective-to-orthographic
                ego_height_flat = tf.reshape(ego_height, [num_images, -1])  # flatten spatial dimensions;
                if self.subtract_min_from_height_map:
                    ego_height_flat -= tf.reduce_min(ego_height_flat)
                if inds_downsamp is not None:
                    assert self.batch_size is not None  # this should never be raised, but just in case;
                    ego_height = tf.gather(ego_height_flat, inds_downsamp, axis=1)  # batch along pixels;
                else:
                    ego_height = ego_height_flat

                if self.use_absolute_scale_calibration:

                    H = self.camera_to_vanish_point_rc[:, None]

                    if self.batch_size is not None:
                        # need to compute the self.j'th camera-to-vanish-point height in case you're batching, which means
                        # that the self.j'th entry may not be computed;
                        n_ground, _ = tf.linalg.normalize(self.ground_surface_normal[self.j], axis=0)
                        n_dot_r = n_ground[2] * self.camera_height[self.j]
                        vanish_xyz = n_dot_r * n_ground  # shape: 3;
                        camera_to_vanish_point_xyz = tf.norm(vanish_xyz, axis=0)
                        H_j = camera_to_vanish_point_xyz * max_dim
                    else:
                        H_j = H[self.j]

                    r = rc_warp - vanish_warp[:, None, :]  # lateral distance to vanishing point;
                    M_j = self.magnification_j
                    f_eff = self.effective_focal_length_mm
                    self.another_height_scale_factor = f_eff * (1 + 1 / M_j)  # scale ego_height again to make this
                    # case similar to the other case (self.use_absolute_scale_calibration);
                    h = ego_height * self.another_height_scale_factor

                    delta_radial = h / f_eff / (1 + 1 / M_j * H / H_j)
                    rc_warp = r * (1 - delta_radial[:, :, None]) + vanish_warp[:, None, :]

                    ego_height *= self.height_scale_factor  # to keep consistent with regularization coefficients;
                    # note that you have to divide by height_scale_factor because the height map is scaled by this, but
                    # you have to divide by another_scale_factor because the multiplication above allows ego_height to
                    # shrink;
                else:
                    ego_height *= self.height_scale_factor  # denominator is large in next line, so multiply by a large
                    # value to allow self.ego_height to take on smaller values;

                    delta_radial = ego_height / self.camera_to_vanish_point_rc[:, None]
                    rc_warp = ((rc_warp - vanish_warp[:, None, :]) * (1 - delta_radial[:, :, None]) +
                               vanish_warp[:, None, :])
                self.ego_height_for_concat = ego_height  # for concatenating with self.im below;
            else:  # perspective-to-perspective warping
                # first, need to define the reference camera view:
                if p2p_warp_mode == 'mean':
                    self.reference_camera_height.assign(tf.reduce_mean(self.camera_to_vanish_point_rc))
                    self.reference_camera_rc.assign(tf.reduce_mean(vanish_warp, axis=0))
                elif p2p_warp_mode == 'random':
                    height_min = tf.reduce_min(self.camera_to_vanish_point_rc)
                    height_max = tf.reduce_max(self.camera_to_vanish_point_rc)
                    self.reference_camera_height.assign(tf.random.uniform((), height_min, height_max))
                    rc_min = tf.reduce_min(vanish_warp, axis=0)
                    rc_max = tf.reduce_max(vanish_warp, axis=0)
                    self.reference_camera_rc.assign(tf.random.uniform((2,), rc_min, rc_max))
                elif p2p_warp_mode == 'fixed':
                    pass  # do nothing, accept current values;
                elif p2p_warp_mode == 'random_choice':
                    # pick one of the camera view among the existing;
                    random_choice = tf.random.uniform((), 0, self.num_images, dtype=tf.int32)
                    self.reference_camera_height.assign(tf.gather(self.camera_to_vanish_point_rc, random_choice))
                    self.reference_camera_rc.assign(tf.gather(vanish_warp, random_choice))
                else:
                    raise Exception('invalid perspective-to-perspective warp mode passed to gradient_update: '
                                    + p2p_warp_mode)
                # vector deformation field to warp to the reference perspective:
                h = tf.reshape(ego_height, [num_images, -1, 1]) * self.height_scale_factor
                H_r = self.reference_camera_height
                H = self.camera_to_vanish_point_rc[:, None, None]
                R_r = self.reference_camera_rc[None, None, :]
                R = vanish_warp[:, None, :]
                r = rc_warp - R  # position vectors relative to each camera's vanishing point;
                p2p_warp = h / (H_r - h) * (R - R_r) + r * h * (H - H_r) / H / (H_r - h)  # the magic equation;
                rc_warp += p2p_warp  # shape: num_images, flattened spatial, 2

                self.ego_height_for_concat = h  # for concatenating with self.im below;

                if self.use_absolute_scale_calibration:
                    raise Exception('not yet implemented for perspective-to-perspective')

        rc_warp = tf.reshape(rc_warp, [-1, 2]) * self.scale  # flatten

        return rc_warp

    def _generate_recon(self, stack_downsamp, rc_downsamp, dither_coords, p2p_warp_mode=None, assign_update_recon=True):
        # backprojects all the images into the reconstruction, with the specified scale;
        # if batching, the gradient_update function will update the reconstruction with a running average;
        # if batching, this function should not be called by the user, as it will continually update the recon with the
        # same batch; if not batching, then this generates the full reconstruction;
        # p2p_warp_mode: if using perspective-to-perspective warping; can be None, 'mean', 'random', or 'fixed';
        # assign_update_recon: only relevant if using batching; controls whether to use the .assign() mechanism to
        # update the reconstruction (specified via update_gradient option in the gradient_update function);

        if self.batch_size is not None:
            if self.batch_across_images:
                # distinguish inds_downsamp and inds_image_downsamp, where the former is for pixel-level batching while
                # the latter is for image-level batching;
                rc_downsamp, inds_image_downsamp = rc_downsamp  # unpack

                # now, for all variables whose first dimension corresponds to the image dimension, gather:
                self.rc_ul_per_im_batch = tf.gather(self.rc_ul_per_im, inds_image_downsamp, axis=0)
                self.gain_batch = tf.gather(self.gain, inds_image_downsamp, axis=0)
                self.bias_batch = tf.gather(self.bias, inds_image_downsamp, axis=0)
                # these are used below:
                gain = self.gain_batch
                bias = self.bias_batch

                if 'camera_parameters' in self.deformation_model:
                    self.camera_height_batch = tf.gather(self.camera_height, inds_image_downsamp, axis=0)
                    self.ground_surface_normal_batch = tf.gather(self.ground_surface_normal,
                                                                 inds_image_downsamp, axis=0)
                    self.camera_in_plane_angle_batch = tf.gather(self.camera_in_plane_angle,
                                                                 inds_image_downsamp, axis=0)
                    if 'unet' not in self.deformation_model:
                        # if using unet, then ego_height will already be gathered as it is generated by the unet;
                        self.ego_height_batch = tf.gather(self.ego_height, inds_image_downsamp, axis=0)
                    # the following self.correct__ variables serve two purposes: 1) to signify whether they are being
                    # used, and 2) specify length of first dimension of the corresponding distortion variable to decide
                    # whether we need to use tf.gather;
                    if self.correct_radial_camera_distortion == self.num_images:
                        self.radial_camera_distortion_batch = tf.gather(self.radial_camera_distortion,
                                                                  inds_image_downsamp, axis=0)
                    if self.correct_radial_camera_distortion_piecewise_linear == self.num_images:
                        self.radial_camera_distortion_piecewise_linear_batch = tf.gather(
                            self.radial_camera_distortion_piecewise_linear, inds_image_downsamp, axis=0)
                    if self.correct_camera_distortion_center == self.num_images:
                        self.camera_distortion_center_batch = tf.gather(self.camera_distortion_center,
                                                                        inds_image_downsamp, axis=0)
                else:
                    raise Exception('image-level batching not yet implemented for a non-camera model')

                inds_downsamp = None
            else:
                # if batching, then stack_downsamp and rc_downsamp are transposed and need to be untransposed;

                if 'camera_parameters' in self.deformation_model:
                    # also need the indices of the pixels chosen, because radial deformations are pixel-wise;
                    # make sure you package these together into a tuple in the script;
                    rc_downsamp, inds_downsamp = rc_downsamp
                else:
                    inds_downsamp = None

                stack_downsamp = tf.transpose(stack_downsamp, (1, 0, 2))
                rc_downsamp = tf.transpose(rc_downsamp, (1, 0, 2))

                # these are used below:
                gain = self.gain
                bias = self.bias
        else:
            inds_downsamp = None
            # these are used below:
            gain = self.gain
            bias = self.bias
        # to save CPU memory, the dataset and coordinates are stored as uint8 and uint16, respectively; thus, cast to
        # float here;
        stack_downsamp = tf.cast(stack_downsamp, self.tf_dtype)
        rc_downsamp = tf.cast(rc_downsamp, self.tf_dtype)

        # function that restricts coordinates to the grid (store as self.variable so that error_map can use it):
        if self.restrict_function == 'clip':
            self.restrict = lambda x: tf.clip_by_value(x, tf.zeros_like(x), self.recon_shape[None] - 1)
        elif 'mod' in self.restrict_function:  # 'mod' or 'mod_with_random_shifts';
            self.restrict = lambda x: tf.math.floormod(x, self.recon_shape[None])
        else:
            raise Exception('invalid restrict_function')

        # apply gain:
        gain_norm = gain / tf.reduce_mean(gain)  # normalize so that there's no global gain;
        im = stack_downsamp * gain_norm[:, None, None] + bias[:, None, None]

        self.im = tf.reshape(im, (-1, self.num_channels))  # flatten all but channels;

        # warped coordinates:
        if self.deformation_model == 'camera_parameters':
            self.rc_warp = self._warp_camera_parameters(rc_downsamp, use_radial_deformation=False)
        elif self.deformation_model == 'camera_parameters_perspective_to_orthographic':
            self.rc_warp = self._warp_camera_parameters(rc_downsamp, inds_downsamp=inds_downsamp,
                                                        use_radial_deformation=True)
        elif self.deformation_model == 'camera_parameters_perspective_to_orthographic_unet':
            self.rc_warp = self._warp_camera_parameters(rc_downsamp, inds_downsamp=inds_downsamp,
                                                        use_radial_deformation=True, stack_downsamp=stack_downsamp)
        elif self.deformation_model == 'camera_parameters_perspective_to_perspective':
            self.rc_warp = self._warp_camera_parameters(rc_downsamp, p2p_warp_mode=p2p_warp_mode,
                                                        inds_downsamp=inds_downsamp,
                                                        use_radial_deformation=True)
        elif self.deformation_model == 'camera_parameters_perspective_to_perspective_unet':
            self.rc_warp = self._warp_camera_parameters(rc_downsamp, p2p_warp_mode=p2p_warp_mode,
                                                        inds_downsamp=inds_downsamp,
                                                        use_radial_deformation=True, stack_downsamp=stack_downsamp)
        else:
            raise Exception('invalid deformation model: ' + self.deformation_model)

        if 'camera_parameters_perspective_to_' in self.deformation_model:
            # adding the height map as a channel to the reconstruction, so first augment self.im with self.ego_height:
            self.im = tf.concat([self.im, tf.reshape(self.ego_height_for_concat, [-1])[:, None]], axis=1)

            self.num_channels_recon = self.num_channels + 1  # for the recon, need one more channel;
        else:
            self.num_channels_recon = self.num_channels

        #
        if dither_coords:
            self.rc_warp += tf.random.uniform([1, 2], -1, 1, dtype=self.tf_dtype)
            if self.batch_size is not None:
                print('Minor warning: using a running average for the recon while dithering coordinates')

        # neighboring pixels:
        rc_floor = tf.floor(self.rc_warp)
        rc_ceil = rc_floor + 1

        # distance to neighboring pixels:
        frc = self.rc_warp - rc_floor
        crc = rc_ceil - self.rc_warp

        # cast
        rc_floor = tf.cast(rc_floor, tf.int32)
        rc_ceil = tf.cast(rc_ceil, tf.int32)

        self.rc_ff = self.restrict(rc_floor)
        self.rc_cc = self.restrict(rc_ceil)
        self.rc_cf = self.restrict(tf.stack([rc_ceil[:, 0], rc_floor[:, 1]], 1))
        self.rc_fc = self.restrict(tf.stack([rc_floor[:, 0], rc_ceil[:, 1]], 1))

        # sig_proj = .42465  # chosen so that if point is exactly in between
        # ...two pixels, .5 weight is assigned to each pixel
        self.frc = tf.exp(-frc ** 2 / 2. / self.sig_proj ** 2)
        self.crc = tf.exp(-crc ** 2 / 2. / self.sig_proj ** 2)

        # augmented coordinates:
        rc_4 = tf.concat([self.rc_ff, self.rc_cc, self.rc_cf, self.rc_fc], 0)

        # interpolated:
        im_4 = tf.concat([self.im * self.frc[:, 0, None] * self.frc[:, 1, None],
                          self.im * self.crc[:, 0, None] * self.crc[:, 1, None],
                          self.im * self.crc[:, 0, None] * self.frc[:, 1, None],
                          self.im * self.frc[:, 0, None] * self.crc[:, 1, None]], 0)
        w_4 = tf.concat([self.frc[:, 0] * self.frc[:, 1],
                         self.crc[:, 0] * self.crc[:, 1],
                         self.crc[:, 0] * self.frc[:, 1],
                         self.frc[:, 0] * self.crc[:, 1]], 0)

        if self.momentum is not None:
            # update with moving average:
            self.im_4_previous = tf.gather_nd(self.recon_previous,
                                              rc_4) * w_4[:, None]  # with appropriate weighting by w_4;
            self.im_4_updated = (im_4 * self.momentum + self.im_4_previous * (1 - self.momentum))
            normalize = tf.scatter_nd(rc_4, w_4, self.recon_shape)
            self.norm_updated_regathered = tf.gather_nd(normalize, rc_4)
            self.im_4_updated_norm = self.im_4_updated / self.norm_updated_regathered[:, None]  # pre-normalize;
            # since tensor_scatter_nd_update doesn't accumulate values, but tensor_scatter_nd_add does, first zero
            # out the regions to be updated and then just add them:
            recon_zeroed = tf.tensor_scatter_nd_update(self.recon_previous, rc_4,
                                                       tf.zeros_like(self.im_4_updated_norm))
            self.recon = tf.tensor_scatter_nd_add(recon_zeroed, rc_4, self.im_4_updated_norm)
            if assign_update_recon:
                with tf.device('/CPU:0'):
                    self.recon_previous.assign(self.recon)

            self.normalize = None  # normalize not needed; in fact, normalize_previous also not needed;

        else:
            self.normalize = tf.scatter_nd(rc_4, w_4, self.recon_shape)
            self.recon = tf.scatter_nd(rc_4, im_4, [self.recon_shape[0], self.recon_shape[1], self.num_channels_recon])
            self.recon = tf.math.divide_no_nan(self.recon, self.normalize[:, :, None])  # creates recon H by W by C;

        if 'camera_parameters_perspective_to' in self.deformation_model:
            self.height_map = self.recon[:, :, -1]
            if self.use_absolute_scale_calibration:
                # divide out the scale factors to get the true height in mm:
                self.tensors_to_track['height_map'] = self.height_map / (self.height_scale_factor /
                                                                         self.another_height_scale_factor)
            else:
                self.tensors_to_track['height_map'] = self.height_map

    def _forward_prediction(self):
        # given the reconstruction, generate forward prediction;

        # forward model:
        ff = tf.gather_nd(self.recon, self.rc_ff)
        cc = tf.gather_nd(self.recon, self.rc_cc)
        cf = tf.gather_nd(self.recon, self.rc_cf)
        fc = tf.gather_nd(self.recon, self.rc_fc)

        self.forward = (ff * self.frc[:, 0, None] * self.frc[:, 1, None] +
                        cc * self.crc[:, 0, None] * self.crc[:, 1, None] +
                        cf * self.crc[:, 0, None] * self.frc[:, 1, None] +
                        fc * self.frc[:, 0, None] * self.crc[:, 1, None])

        self.forward /= ((self.frc[:, 0, None] * self.frc[:, 1, None]) +
                         (self.crc[:, 0, None] * self.crc[:, 1, None]) +
                         (self.crc[:, 0, None] * self.frc[:, 1, None]) +
                         (self.frc[:, 0, None] * self.crc[:, 1, None]))

        if 'camera_parameters_perspective' in self.deformation_model:
            # split off the last dimension, the height dimension, to compute the height map MSE:
            self.forward_height = self.forward[:, -1]
            self.error_height = self.forward_height - self.im[:, -1]  # save this for computing error map;
            self.MSE_height = tf.reduce_mean(self.error_height ** 2)

            self.error = self.forward[:, :-1] - self.im[:, :-1]  # remaining channels are the actual recon;
            self.MSE = tf.reduce_mean(self.error ** 2)
            self.recon = self.recon[:, :, :-1]  # discard the height map channel, as it's recorded elsewhere;
        else:
            self.error = self.forward - self.im  # save this for computing error map;
            self.MSE = tf.reduce_mean(self.error ** 2)

        if self.report_error_map:
            # project the squared error onto the reconstruction space;
            assert self.momentum is None  # for now, don't use this with batching/momentum;
            # don't have to use interpolated projection, because we don't need to compute gradients for this;

            with tf.device('/CPU:0'):
                rc_warp_int = self.restrict(tf.cast(self.rc_warp, tf.int32))
                self.error_map = tf.scatter_nd(rc_warp_int, self.error ** 2,
                                               [self.recon_shape[0], self.recon_shape[1], self.num_channels])
                norm = tf.scatter_nd(rc_warp_int, tf.ones(self.rc_warp.shape[0], dtype=self.tf_dtype), self.recon_shape)
                self.error_map = tf.math.divide_no_nan(self.error_map, norm[:, :, None])
        else:
            self.error_map = None

    def _add_regularization_loss(self, reg_coefs):

        # tf Variable to be regularized:
        if 'perspective_to_' in self.deformation_model:
            field = self.ego_height_to_regularize

        self.loss_list = [self.MSE]
        if 'L2' in reg_coefs and reg_coefs['L2'] is not None:
            self.loss_list.append(reg_coefs['L2'] * tf.reduce_sum(field ** 2) / self.scale ** 2)
        if 'L1' in reg_coefs and reg_coefs['L1'] is not None:
            self.loss_list.append(reg_coefs['L1'] * tf.reduce_sum(tf.sqrt(field ** 2 + 1e-7)) / self.scale ** 2)
        if 'TV2' in reg_coefs and reg_coefs['TV2'] is not None:  # TV squared;
            d0 = field[:, 1:, :-1] - field[:, :-1, :-1]
            d1 = field[:, :-1, 1:] - field[:, :-1, :-1]
            self.loss_list.append(reg_coefs['TV2'] * tf.reduce_sum(d0 ** 2 + d1 ** 2) / self.scale ** 2)
        if 'TV' in reg_coefs and reg_coefs['TV'] is not None:
            d0 = field[:, 1:, :-1] - field[:, :-1, :-1]
            d1 = field[:, :-1, 1:] - field[:, :-1, :-1]
            self.loss_list.append(reg_coefs['TV'] * tf.reduce_sum(tf.sqrt(d0 ** 2 + d1 ** 2 + 1e-7)) / self.scale ** 2)
        if 'height_map' in reg_coefs and reg_coefs['height_map'] is not None:
            # regularize consistency of deformations to the global height map;
            self.loss_list.append(reg_coefs['height_map'] * self.MSE_height)

    @tf.function
    def gradient_update(self, stack_downsamp, rc_downsamp, update_gradient=True, reg_coefs=None, dither_coords=True,
                        return_tracked_tensors=False, stop_gradient=True, p2p_warp_mode=None, return_grads=False,
                        return_loss_only=False):
        # do one gradient update step; pass thru the downsampled stack and rc coordinates via tf.datasets if batching;
        # one might set update_gradient to False if using batching and just want to create the reconstruction without
        # registering simultaneously;
        # dither_coords: if True, add a random float between -1 and 1 to prevent adaptation to the interpolation scheme;
        # stop_gradient: whether to use tf.stop_gradient on the reconstruction before computing the forward prediction;
        # p2p_warp_mode: see warp_camera_parameters and/or generate_recon; if using perspective-to-orthographic warping,
        # then set this to None; if using perspective-to-perspective warping, then set to one of 'mean', 'random',
        # 'fixed';
        # return_loss_only: as the name suggests;

        with tf.GradientTape() as tape:
            self._generate_recon(stack_downsamp, rc_downsamp, dither_coords, p2p_warp_mode, update_gradient)
            if stop_gradient:
                self.recon = tf.stop_gradient(self.recon)
            self._forward_prediction()

            if reg_coefs is not None:
                self._add_regularization_loss(reg_coefs)
                loss = tf.reduce_sum(self.loss_list)
            else:
                loss = self.MSE

        grads = tape.gradient(loss, self.train_var_list)
        # apply gradient update for each optimizer:
        if update_gradient:
            for grad, var, optimizer, train in zip(grads, self.train_var_list, self.optimizer_list,
                                                   self.trainable_or_not):
                if train:  # if user specifies negative learning rate, then don't update;
                    if type(var) is list or type(var) is ListWrapper:  # sometimes it's a ListWrapper; after restoring
                        # from checkpoint?
                        # this is probably the neural network variable list;
                        optimizer.apply_gradients(zip(grad, var))
                    else:
                        optimizer.apply_gradients([(grad, var)])

            # update other stuff:
            if 'camera_parameters' in self.deformation_model and self.force_ground_surface_up and update_gradient:
                # rotate all vectors such that their mean direction points in [0,0,-1];
                # (mean across all, not just batch!);
                # using rodrigues formula;
                n_mean = tf.reduce_mean(self.ground_surface_normal, axis=0)
                rot_axis = tf.stack([-n_mean[1], n_mean[0], 0])  # cross product;
                rot_axis, _ = tf.linalg.normalize(rot_axis)
                cos_rot_angle = -n_mean[2] / tf.norm(n_mean)
                sin_rot_angle = tf.sqrt(1 - cos_rot_angle ** 2 + 1e-6)  # prevent numerical issues;
                K = tf.stack([[0, -rot_axis[2], rot_axis[1]],
                              [rot_axis[2], 0, -rot_axis[0]],
                              [-rot_axis[1], rot_axis[0], 0]])
                rotmat = tf.eye(3) + sin_rot_angle * K + K @ K * (1 - cos_rot_angle)
                self.ground_surface_normal.assign(tf.einsum('ij,cj->ci', rotmat, self.ground_surface_normal))

        if return_loss_only:
            if reg_coefs is not None:
                return self.loss_list
            else:
                return self.MSE
        else:
            if reg_coefs is not None:
                return_list = [self.loss_list, self.recon, self.normalize, self.error_map]
            else:
                return_list = [self.MSE, self.recon, self.normalize, self.error_map]

            if return_tracked_tensors:
                return_list.append(self.tensors_to_track)
            if return_grads:
                return_list.append(grads)

            return return_list

    def checkpoint_all_variables(self, path='./tf_ckpts', skip_saving=False):
        # save checkpoints to restore in case unet diverges;
        # skip_saving if you just want to create the ckpt and manager;
        if self.ckpt is None:
            self.ckpt = tf.train.Checkpoint()
            self.ckpt.var = self.train_var_list
            self.ckpt.opt = self.optimizer_list
            self.ckpt.non_train_var = self.non_train_list
            self.manager = tf.train.CheckpointManager(self.ckpt, path, max_to_keep=2)
            # only keep two, restore the oldest;
        if not skip_saving:
            self.manager.save()

    def restore_all_variables(self):
        self.ckpt.restore(self.manager.checkpoints[0])

    @tf.function
    def get_all_variables(self):
        # return as dictionary so it can be used as initialization:
        with tf.device('/CPU:0'):
            # don't return neural network parameters -- use checkpoints instead;
            vars = dict()
            for var in self.train_var_list:
                try:
                    vars[var.name[:-2]] = var
                except AttributeError:  # ListWrapper for unet, for which we don't need the variables;
                    pass
            return vars


# u-net for mapping images to height map:
class unet(tf.keras.Model):
    def __init__(self, filters_list, skip_list, output_nonlinearity, upsample_method='bilinear'):
        # filters_list and skip_list are lists of number of filters in the upsample/downsample layers,
        # and the number of filters in the skip connections;
        # output_nonlinearity can be 'leaky_relu' or 'linear';
        # upsample_method can be 'bilinear' or 'nearest', which is used in the upsample blocks;
        super(unet, self).__init__()
        assert len(filters_list) == len(skip_list)
        self.filters_list = filters_list
        self.skip_list = skip_list
        self.output_nonlinearity = output_nonlinearity
        self.upsample_method = upsample_method
        self._build()

    def _build(self):
        # define all the layers of the unet;
        self.downsample_list = list()  # stores list of downsample blocks;
        self.downsample_skip_block_list = list()  # stores list of skip convolutional blocks;
        self.upsample_list = list()  # stores list of upsample blocks;
        self.upsample_concat_list = list()  # stores list of concatenation layers;

        # downsampling half:
        for num_filters, num_skip_filters in zip(self.filters_list, self.skip_list):
            self.downsample_list.append(self._downsample_block(num_filters))  # add to list of layers
            self.downsample_skip_block_list.append(self._skip_block(num_skip_filters))

        # upsampling half:
        for i, (num_filters, num_skip_filters) in enumerate(zip(self.filters_list[::-1], self.skip_list[::-1])):
            if num_skip_filters != 0:
                self.upsample_concat_list.append(tf.keras.layers.Concatenate())
            else:
                self.upsample_concat_list.append(None)  # as a placeholder
            if i == len(self.filters_list) - 1:
                # last block, use the specified output nonlinearity:
                self.upsample_list.append(self._upsample_block(num_filters,
                                                               nonlinearity=self.output_nonlinearity))
            else:
                self.upsample_list.append(self._upsample_block(num_filters))

    def _downsample_block(self, numfilters, kernel_size=3):
        return [tf.keras.layers.Conv2D(filters=numfilters, kernel_size=kernel_size,
                                       strides=(2, 2), padding='same'),
                tf.keras.layers.BatchNormalization(),
                tf.keras.layers.LeakyReLU(),
                # repeat, but no downsample this time
                tf.keras.layers.Conv2D(filters=numfilters, kernel_size=3,
                                       strides=(1, 1), padding='same'),
                tf.keras.layers.BatchNormalization(),
                tf.keras.layers.LeakyReLU()]

    def _upsample_block(self, numfilters, kernel_size=3, nonlinearity='leaky_relu'):
        layers_list = [tf.keras.layers.UpSampling2D(interpolation=self.upsample_method),
                       tf.keras.layers.Conv2D(filters=numfilters, kernel_size=kernel_size,
                                              strides=(1, 1), padding='same'),
                       tf.keras.layers.BatchNormalization(),
                       tf.keras.layers.LeakyReLU(),
                       tf.keras.layers.Conv2D(filters=numfilters, kernel_size=1,
                                              strides=(1, 1), padding='same'),  # kernel size 1
                       tf.keras.layers.BatchNormalization()]
        if nonlinearity == 'leaky_relu':
            layers_list.append(tf.keras.layers.LeakyReLU())
        elif nonlinearity == 'linear':
            pass
        else:
            raise Exception('invalid nonlinearity')
        return layers_list

    def _skip_block(self, numfilters=4, kernel_size=1):
        if numfilters == 0:  # no skip connections
            return None
        else:
            return [tf.keras.layers.Conv2D(filters=numfilters, kernel_size=kernel_size,
                                           strides=(1, 1), padding='same'),
                    tf.keras.layers.BatchNormalization(),
                    tf.keras.layers.LeakyReLU()]

    def call(self, x):
        skip_layers = list()  # store skip layer outputs to be concatenated in the upsample block;
        for down_block, skip_block in zip(self.downsample_list, self.downsample_skip_block_list):
            for down_layer in down_block:  # traverse all layers in block;
                x = down_layer(x)
            if skip_block is not None:  # if there's a skip block, traverse all layers in it;
                x_ = x
                for skip_layer in skip_block:
                    x_ = skip_layer(x_)
                skip_layers.append(x_)
            else:
                skip_layers.append(None)
        for up_block, skip, concat in zip(self.upsample_list, skip_layers[::-1], self.upsample_concat_list):
            if skip is not None:
                x = concat([x, skip])
            for up_layer in up_block:
                x = up_layer(x)
        return x


# this function is used with the unet to calculate how to pad the input (and subsequently depad the output):
def get_compatible_size(dim, num_downsamp, max_dim=10000):
    # for a given dim and number of downsample blocks, find the smallest value >= dim such that the
    # unet will return the same size as the input;
    # max_dim is the largest value to be considered;
    k = 2 ** num_downsamp
    for i in range(1, max_dim):
        new_dim = i * k
        if new_dim >= dim:
            break
    return new_dim

# below are convenience functions for the script to call:
def stack_loader_phone(directory):
    # for my phone datasets;
    im_stack = list()
    for filename in tqdm(sorted(os.listdir(directory))):
        if 'jpg' in filename or 'png' in filename:
            im = plt.imread(directory + filename)
            im_stack.append(im)
    im_stack = np.array(im_stack)
    shape = im_stack.shape[1:]
    print('Shape of stack: ' + str(shape))
    return im_stack

def xcorr_initial_guess(im_stack, downsamp=4, channel=0, crop_frac=.4):
    # generates initial guess for x_s and y_pos using successive cross correlations (therefore, sequential images should
    # have overlap);
    # expect the input im_stack to be the output of an image loading script;
    # specify downsample factor (without antialiasing) and which channel to use;
    # crop_frac -- only search a center crop of the xcorr;

    im_stack_ = im_stack[:, ::downsamp, ::downsamp, channel].astype(np.float32)
    pos = [[0, 0]]  # first coordinate is arbitrarily defined to be the origin;
    for i in range(len(im_stack) - 1):
        xcorr = scipy.signal.correlate(im_stack_[i], im_stack_[i + 1])
        bkgd = scipy.signal.correlate(im_stack_[i], im_stack_[i])  # use autocorrelation as background;
        xcorr -= bkgd  # subtract background to avoid bias approaching 0 delay
        s = np.array(xcorr.shape)
        start_pix = np.int32((s - s * crop_frac) / 2)
        end_pix = np.int32((s + s * crop_frac) / 2)
        xcorr = xcorr[start_pix[0]:end_pix[0], start_pix[1]:end_pix[1]]

        argmax = np.unravel_index(xcorr.argmax(), xcorr.shape)
        pos.append(argmax - np.array(xcorr.shape) / 2)
    pos = np.array(pos).astype(np.float32) * downsamp
    pos = np.cumsum(pos, axis=0)  # convert relative shifts into absolute positions
    x_pos = pos[:, 1]
    y_pos = pos[:, 0]
    x_pos -= np.min(x_pos)  # avoid negative values because they will be stored as uints
    y_pos -= np.min(y_pos)

    return x_pos, y_pos

def monitor_progress(recon, error_map, losses, tracked=None):
    # run this in your train loop to print out progress;
    # recon is the reconstruction, error_map is the error map, losses is a list of MSEs or loss terms, and tracked
    # is a dictionary of tracked tensors from the tf graph;

    num_channels = recon.shape[-1]
    plt.figure(figsize=(15, 15))
    if num_channels == 1:
        plt.imshow(recon[:, :, 0].numpy(), cmap='jet')
        plt.colorbar()
    else:
        plt.imshow(np.uint8(np.clip(recon.numpy(), 0, 255)))
    plt.title('Reconstruction')
    plt.show()

    if error_map is not None:
        plt.figure(figsize=(15, 15))
        plt.imshow(error_map.numpy().sum(2))
        plt.clim([0, 300])
        plt.colorbar()
        plt.title('Error map')
        plt.show()

    if tracked is not None and 'height_map' in tracked:
        height_map = tracked['height_map']
        mask = recon.numpy().sum(2) > 0  # which pixels were visited (also could use normalize);
        clims = np.percentile(height_map[mask], [.1, 99.9])
        plt.figure(figsize=(15, 15))
        plt.imshow(height_map, cmap='jet')
        plt.clim(clims)
        plt.title('Height map')
        plt.colorbar()
        plt.show()

    plt.plot(losses)
    plt.title('Loss history')
    plt.show()
# stack: stack of images to be stitched; 3D tensor of shape num_images, num_rows, num_cols, num_channels; # ul_coords: array of upper-left coordinates of the images (in pixels); 2D tensor of shape num_images, 2 (row, # col; y, x); # recon_shape: shape of the final reconstruction/stitched image (row, col, channels); # ul_offset: length-2 vector (row, col) specifying a constant offset for all the ul_coords (in pixels); # batch_size: if None, then don't use batching; # scale: factor between 0 and 1 specifying how much to downsample the reconstruction; # momentum: only relevant for batching; specifies how much averaging of previous iterations to use for recon; # either both batch_size and momentum must be non-None, or both must be None; # batch_across_images: if you're batching (i.e., batch_size is not None), then if true, then batch across the # image dimension; otherwise, batch across pixels; (if batching while using unet, then you must batch across # image dimension); # report_error_map: # restrict_function choice decides what to do if a point goes beyond the boundaries; 'clip' or 'mod' or # 'mod_with_random_shifts'; # cast is 8-bit to save memory for the batch generator; cast batch to float32; # number of channels; stack must at least have a singleton dim 3; # number of images in dataset; # uint8 is too narrow, because these pixel coordinates are large; # this will be the base recon_shape; the effective recon_shape will depend # on the scale factor that the user specifies; # for the intepolation kernel width; # unet parameters if relevant; define these manually if needed; # to scale the output of the unet # 'linear' or 'leaky_relu' # 'bilinear' or 'nearest' # for checkpointing models when using unet; # save model every this many iterations; # save memory using tf.recompute_grad; # to scale the ego_height parameter to avoid dealing with large values; # camera parameters for getting absolute scale; # effective focal length in mm; # user needs to figure out the magnification of the jth camera view; # and specify which camera view index; by default, assume the first camera; # define tf.Variables and optimizers; # deformation_model: affects what tf.Variables will be defined and optimized; # learning_rates is a dictionary of variable names (strings) and the corresponding learning rate; None means # use the default, defined below; can specify a learning rate a negative number, indicating that variable is not # optimized; can supply initial variable values as a dictionary; if None, use default; # if recon and normalize are supplied, then initialize recon_previous and normalize_previous with these after # upsampling; both must be supplied; (this is for momentum/running average calculations); most likely, recon and # normalize will be from earlier reconstruction attempts using mesoSfM, but at lower resolution; # remove_global_transform: e.g., no global shift or scale; # stack_downsample_factor is an integer that downsamples the stack and coordinates; if None, then it will be # computed from self.scale; # force_ground_surface_up; when using camera model, force the mean surface normal to be [0,0,-1]; # these must be both be None, or neither None: # define downsample factor: # also downsample the images to save computation; # obviously can't downsample with 0; # coordinates of the images: # base coordinates (could be downsampled after applying a # scale factor); # only relevant if using a camera model; # list of tf.Variables and optimizers, to be populated by one or more _create_variables_*deformation_model*; # list of variables that aren't trained (probably .assign()'d; for checkpoints); # intermediate tensors to track; have a tf.function return the contents; # to be used directly below; allows for dictionaries in which not all keys are specified; if not specified, # then use default_dict's value; # if nothing given, then use the default; # if the key is present, but None is specified; # i.e., use the value given; # if key is not even present; # user has to manually define these: # initialize same value as ... # ...for height; # to avoid /0; # add small value to allow gradients to prop; # these are for perspective to perspective and perspective to orthographic: # create radial deformation field: # create neural network layers: # removing the global scaling transform entails adapting the focal length to the mean height; # these won't be optimized; also these initializing values are very temporary, and should be immediately # modified by code later; the reason why they are tf.Variables is so that the user can manually modify # them in eager mode with .assign(); # note that because there's no tilt, reference_camera_height is the same as the camera to vanishing # point distance, and the reference_camera_rc is the same as the vanishing point position; # intensity adjustment (uniform within a given image for now): # intensity bias (also uniform within each image): # barrel or pincushion distortion correction: # this is explicitly a camera modeling option; # user is specifying there to be only one common parameter; # make sure at least 2 dims; # assume this one dimension refers to the camera view, not the polynomial order; # second dim: polynomial order; # number of terms in the polynomial (even only); # half of the even powers to raise to; # this info used if batching; # for now, only allow a common distortion among all cams; # length of this vector determines how many discretization levels; # how many discretization levels (nodes); # used if batching images; this never equals # num_images; # if camera is not centered; this defines the center of the above two distortions, so center the camera first, # apply the distortions, then decenter back; # one common pair of parameters; # x and y centers; # by default, diff pair for each camera; # this info used if # batching; # create a list of booleans to accompany self.train_var_list and self.optimizer_list to specify whether to train # those variables (as specified by the whether the user-specified learning rates are negative); doing this so # that autograph doesn't traverse all branches of the conditionals; if the user ever wants to turn off # optimization of a variable mid-optimization, then just do .assign(0) to the learning rate, such that the # update is still happening, but the change is 0; # if the variable is a list of variables, then this should be for the unet; modify here if there are # other scenarios; # downsample rc coordinates and stack: # cv2 doesn't like numpy values? # must be odd; # flatten spatial dims; # flatten spatial dims; # create variables relevant for batching (or set them to None if not batching): # we're going to give the coregistered height map a ride; # if none supplied, initialize with 0s; # otherwise, upsample to the current shape; # cv2 seems to squeeze singleton channels dimensions, so if it's singleton, add it back: # if you would like to use the previous RGB image as initialization, but your previous run didn't # estimate a height map; # add empty height channel: # initialize first recon and normalize tensors for momentum; use the scaled recon shape, not the base shape; # radially inwardly pointing vector magnitudes, where the larger the magnitude, the taller the object; # make sure the first dimensions match (number of images in stack): # presumably you've initialized this with the results from another optimization at a different scale; # thus, resize to match the current scale: # convert from tf to np if you need to: # run the network once so that we can access network.trainable_variables # get padded shape that the network likes; # coordinates are in camera space; # height from the camera perspective; # user calls this function to get a dataset to iterate over; if not using batching, then just return a tuple or # list of length 1 (i.e., the whole dataset is one batch); # sample a subset of the images, and keep track of the indices downsampled so that you can gather the # corresponding variables; # transpose to batch along space, not image number; # need to also get coordinates of the spatial positions to index into pixel-wise deformation fields: # basically a 1-batch dataset; # shape of rc_downsamp: num_images, _, 2; # use_radial_deformation is a boolean flag specifying whether to do the per-pixel radial deformation fields to # warp perspective to othographic OR perspective to perspective; if the latter, then p2p_warp_mode specifies # how to specify the perspective reference to warp to; the options are 'mean', 'random', 'fixed', and None, # where None means you're using perspective-to-orthographpic warping and not perspective to perspective; # inds_downsamp is passed if batching and using radial deformations; # stack_downsamp is only needed if using a unet; # in generate_recon, we defined the batch and non-batch versions; # for reshaping below; # not batched; # for reshaping below; # don't use self.camera_focal_length; set to geometric mean of the camera heights; # also, use the tf.Variable version always (not the batch version); # for normalization of image coordinates; # to keep isotropic; # normalize to -.5 to .5; # radial distance squared; # dims^: camera, pixels # make it go from -1 to 1 rather than -.5 to .5; # even polynomial to account for distortion: # dims^: camera, pixels, power # dims^: camera, pixels # radial distance; # dims^: camera, pixels; these radial coordinates should be between 0 and .5*sqrt(2), but could go higher # if th center moves); thus to be safe, just multiply by num_radial_pixels; # find nearest pixels and distances thereto: # to prevent out of range indexing; # to prevent global expansion; # bilinear interpolation: # in-plane rotation: # normalize to unit mag; # shape^: num_images, 3 # projecting to object space (computed analytically and taylor-expanded); # using a taylor expansion: # shape: num_images; dot product between a point on the # ground, r (0, 0, camera_height); (this is needed below); # compute the vanishing point, which will be used if you use the projective to orthographic mode; # shape: num_images, 3; # distance from camera to the ground; the actual # height of the camera; # projection to object space simplifies to this: # vanishing point in camera plane: # follow surface normal to camera # plane; # don't need z, it's just the focal length; # account for in-plane camera rotation: # convert back to row-column: # convert from xy units to rc units; # add translations (same as for the homography and affine implementations): # these translations don't affect camera_to_vanish_point; # to discourage registration with overlapped regions; # generate self.ego_height; doesn't matter if batching or not because it's generated from the image batch; # remove last dimension; # need this for regularization; # multiplicative version: # using perspective-to-orthographic # flatten spatial dimensions; # this should never be raised, but just in case; # batch along pixels; # need to compute the self.j'th camera-to-vanish-point height in case you're batching, which means # that the self.j'th entry may not be computed; # shape: 3; # lateral distance to vanishing point; # scale ego_height again to make this # case similar to the other case (self.use_absolute_scale_calibration); # to keep consistent with regularization coefficients; # note that you have to divide by height_scale_factor because the height map is scaled by this, but # you have to divide by another_scale_factor because the multiplication above allows ego_height to # shrink; # denominator is large in next line, so multiply by a large # value to allow self.ego_height to take on smaller values; # for concatenating with self.im below; # perspective-to-perspective warping # first, need to define the reference camera view: # do nothing, accept current values; # pick one of the camera view among the existing; # vector deformation field to warp to the reference perspective: # position vectors relative to each camera's vanishing point; # the magic equation; # shape: num_images, flattened spatial, 2 # for concatenating with self.im below; # flatten # backprojects all the images into the reconstruction, with the specified scale; # if batching, the gradient_update function will update the reconstruction with a running average; # if batching, this function should not be called by the user, as it will continually update the recon with the # same batch; if not batching, then this generates the full reconstruction; # p2p_warp_mode: if using perspective-to-perspective warping; can be None, 'mean', 'random', or 'fixed'; # assign_update_recon: only relevant if using batching; controls whether to use the .assign() mechanism to # update the reconstruction (specified via update_gradient option in the gradient_update function); # distinguish inds_downsamp and inds_image_downsamp, where the former is for pixel-level batching while # the latter is for image-level batching; # unpack # now, for all variables whose first dimension corresponds to the image dimension, gather: # these are used below: # if using unet, then ego_height will already be gathered as it is generated by the unet; # the following self.correct__ variables serve two purposes: 1) to signify whether they are being # used, and 2) specify length of first dimension of the corresponding distortion variable to decide # whether we need to use tf.gather; # if batching, then stack_downsamp and rc_downsamp are transposed and need to be untransposed; # also need the indices of the pixels chosen, because radial deformations are pixel-wise; # make sure you package these together into a tuple in the script; # these are used below: # these are used below: # to save CPU memory, the dataset and coordinates are stored as uint8 and uint16, respectively; thus, cast to # float here; # function that restricts coordinates to the grid (store as self.variable so that error_map can use it): # 'mod' or 'mod_with_random_shifts'; # apply gain: # normalize so that there's no global gain; # flatten all but channels; # warped coordinates: # adding the height map as a channel to the reconstruction, so first augment self.im with self.ego_height: # for the recon, need one more channel; # # neighboring pixels: # distance to neighboring pixels: # cast # sig_proj = .42465  # chosen so that if point is exactly in between # ...two pixels, .5 weight is assigned to each pixel # augmented coordinates: # interpolated: # update with moving average: # with appropriate weighting by w_4; # pre-normalize; # since tensor_scatter_nd_update doesn't accumulate values, but tensor_scatter_nd_add does, first zero # out the regions to be updated and then just add them: # normalize not needed; in fact, normalize_previous also not needed; # creates recon H by W by C; # divide out the scale factors to get the true height in mm: # given the reconstruction, generate forward prediction; # forward model: # split off the last dimension, the height dimension, to compute the height map MSE: # save this for computing error map; # remaining channels are the actual recon; # discard the height map channel, as it's recorded elsewhere; # save this for computing error map; # project the squared error onto the reconstruction space; # for now, don't use this with batching/momentum; # don't have to use interpolated projection, because we don't need to compute gradients for this; # tf Variable to be regularized: # TV squared; # regularize consistency of deformations to the global height map; # do one gradient update step; pass thru the downsampled stack and rc coordinates via tf.datasets if batching; # one might set update_gradient to False if using batching and just want to create the reconstruction without # registering simultaneously; # dither_coords: if True, add a random float between -1 and 1 to prevent adaptation to the interpolation scheme; # stop_gradient: whether to use tf.stop_gradient on the reconstruction before computing the forward prediction; # p2p_warp_mode: see warp_camera_parameters and/or generate_recon; if using perspective-to-orthographic warping, # then set this to None; if using perspective-to-perspective warping, then set to one of 'mean', 'random', # 'fixed'; # return_loss_only: as the name suggests; # apply gradient update for each optimizer: # if user specifies negative learning rate, then don't update; # sometimes it's a ListWrapper; after restoring # from checkpoint? # this is probably the neural network variable list; # update other stuff: # rotate all vectors such that their mean direction points in [0,0,-1]; # (mean across all, not just batch!); # using rodrigues formula; # cross product; # prevent numerical issues; # save checkpoints to restore in case unet diverges; # skip_saving if you just want to create the ckpt and manager; # only keep two, restore the oldest; # return as dictionary so it can be used as initialization: # don't return neural network parameters -- use checkpoints instead; # ListWrapper for unet, for which we don't need the variables; # u-net for mapping images to height map: # filters_list and skip_list are lists of number of filters in the upsample/downsample layers, # and the number of filters in the skip connections; # output_nonlinearity can be 'leaky_relu' or 'linear'; # upsample_method can be 'bilinear' or 'nearest', which is used in the upsample blocks; # define all the layers of the unet; # stores list of downsample blocks; # stores list of skip convolutional blocks; # stores list of upsample blocks; # stores list of concatenation layers; # downsampling half: # add to list of layers # upsampling half: # as a placeholder # last block, use the specified output nonlinearity: # repeat, but no downsample this time # kernel size 1 # no skip connections # store skip layer outputs to be concatenated in the upsample block; # traverse all layers in block; # if there's a skip block, traverse all layers in it; # this function is used with the unet to calculate how to pad the input (and subsequently depad the output): # for a given dim and number of downsample blocks, find the smallest value >= dim such that the # unet will return the same size as the input; # max_dim is the largest value to be considered; # below are convenience functions for the script to call: # for my phone datasets; # generates initial guess for x_s and y_pos using successive cross correlations (therefore, sequential images should # have overlap); # expect the input im_stack to be the output of an image loading script; # specify downsample factor (without antialiasing) and which channel to use; # crop_frac -- only search a center crop of the xcorr; # first coordinate is arbitrarily defined to be the origin; # use autocorrelation as background; # subtract background to avoid bias approaching 0 delay # convert relative shifts into absolute positions # avoid negative values because they will be stored as uints # run this in your train loop to print out progress; # recon is the reconstruction, error_map is the error map, losses is a list of MSEs or loss terms, and tracked # is a dictionary of tracked tensors from the tf graph; # which pixels were visited (also could use normalize);