def _build_egomotion_test_graph(self): """Builds graph for inference of egomotion given two images.""" with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE): self._image1 = tf.placeholder( tf.float32, [self.batch_size, self.img_height, self.img_width, 3], name='image1') self._image2 = tf.placeholder( tf.float32, [self.batch_size, self.img_height, self.img_width, 3], name='image2') # The "compute_loss" scope is needed for the checkpoint to load properly. with tf.name_scope('compute_loss'): rot, trans, _, _ = motion_prediction_net.motion_field_net( images=tf.concat([self._image1, self._image2], axis=-1)) inv_rot, inv_trans, _, _ = ( motion_prediction_net.motion_field_net(images=tf.concat( [self._image2, self._image1], axis=-1))) rot = transform_utils.matrix_from_angles(rot) inv_rot = transform_utils.matrix_from_angles(inv_rot) trans = tf.squeeze(trans, axis=(1, 2)) inv_trans = tf.squeeze(inv_trans, axis=(1, 2)) # rot and inv_rot should be the inverses on of the other, but in reality # they slightly differ. Averaging rot and inv(inv_rot) gives a better # estimator for the rotation. Similarly, trans and rot*inv_trans should # be the negatives one of the other, so we average rot*inv_trans and trans # to get a better estimator. TODO(gariel): Check if there's an estimator # with less variance. self.rot = 0.5 * (tf.linalg.inv(inv_rot) + rot) self.trans = 0.5 * (-tf.squeeze( tf.matmul(self.rot, tf.expand_dims(inv_trans, -1)), axis=-1) + trans)
def _build_loss(self): """Builds the loss tensor, to be minimized by the optimizer.""" self.reader = reader.DataReader( self.data_dir, self.batch_size, self.img_height, self.img_width, SEQ_LENGTH, 1, # num_scales self.file_extension, self.random_scale_crop, reader.FLIP_RANDOM, self.random_color, self.imagenet_norm, self.shuffle, self.input_file, queue_size=self.queue_size) (self.image_stack, self.image_stack_norm, self.seg_stack, self.intrinsic_mat, _) = self.reader.read_data() if self.learn_intrinsics: self.intrinsic_mat = None if self.intrinsic_mat is None and not self.learn_intrinsics: raise RuntimeError( 'Could not read intrinsic matrix. Turn ' 'learn_intrinsics on to learn it instead of loading ' 'it.') self.export('self.image_stack', self.image_stack) object_masks = [] for i in range(self.batch_size): object_ids = tf.unique(tf.reshape(self.seg_stack[i], [-1]))[0] object_masks_i = [] for j in range(SEQ_LENGTH): current_seg = self.seg_stack[i, :, :, j * 3] # (H, W) def process_obj_mask(obj_id): """Create a mask for obj_id, skipping the background mask.""" mask = tf.logical_and( tf.equal(current_seg, obj_id), # pylint: disable=cell-var-from-loop tf.not_equal(tf.cast(0, tf.uint8), obj_id)) # Leave out vert small masks, that are most often errors. size = tf.reduce_sum(tf.to_int32(mask)) mask = tf.logical_and(mask, tf.greater(size, MIN_OBJECT_AREA)) if not self.boxify: return mask # Complete the mask to its bounding box. binary_obj_masks_y = tf.reduce_any(mask, axis=1, keepdims=True) binary_obj_masks_x = tf.reduce_any(mask, axis=0, keepdims=True) return tf.logical_and(binary_obj_masks_y, binary_obj_masks_x) object_mask = tf.map_fn( # (N, H, W) process_obj_mask, object_ids, dtype=tf.bool) object_mask = tf.reduce_any(object_mask, axis=0) object_masks_i.append(object_mask) object_masks.append(tf.stack(object_masks_i, axis=-1)) self.seg_stack = tf.cast(tf.stack(object_masks, axis=0), tf.float) tf.summary.image('Masks', self.seg_stack) with tf.variable_scope(DEPTH_SCOPE): # Organized by ...[i][scale]. Note that the order is flipped in # variables in build_loss() below. self.disp = {} self.depth = {} # Parabolic rampup of he noise over LAYER_NORM_NOISE_RAMPUP_STEPS steps. # We stop at 0.5 because this is the value above which the multiplicative # noise we use can become negative. Further experimentation is needed to # find if non-negativity is indeed needed. noise_stddev = 0.5 * tf.square( tf.minimum( tf.cast(self.global_step, tf.float) / float(LAYER_NORM_NOISE_RAMPUP_STEPS), 1.0)) def _normalizer_fn(x, is_train, name='bn'): return randomized_layer_normalization.normalize( x, is_train=is_train, name=name, stddev=noise_stddev) with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE): for i in range(SEQ_LENGTH): image = self.image_stack_norm[:, :, :, 3 * i:3 * (i + 1)] self.depth[ i] = depth_prediction_net.depth_prediction_resnet18unet( image, True, self.weight_reg, _normalizer_fn) self.disp[i] = 1.0 / self.depth[i] with tf.name_scope('compute_loss'): self.reconstr_loss = 0 self.smooth_loss = 0 self.ssim_loss = 0 self.depth_consistency_loss = 0 # Smoothness. if self.smooth_weight > 0: for i in range(SEQ_LENGTH): disp_smoothing = self.disp[i] # Perform depth normalization, dividing by the mean. mean_disp = tf.reduce_mean(disp_smoothing, axis=[1, 2, 3], keep_dims=True) disp_input = disp_smoothing / mean_disp self.smooth_loss += _depth_smoothness( disp_input, self.image_stack[:, :, :, 3 * i:3 * (i + 1)]) self.rot_loss = 0.0 self.trans_loss = 0.0 def add_result_to_loss_and_summaries(endpoints, i, j): tf.summary.image( 'valid_mask%d%d' % (i, j), tf.expand_dims(endpoints['depth_proximity_weight'], -1)) self.depth_consistency_loss += endpoints['depth_error'] self.reconstr_loss += endpoints['rgb_error'] self.ssim_loss += 0.5 * endpoints['ssim_error'] self.rot_loss += endpoints['rotation_error'] self.trans_loss += endpoints['translation_error'] self.motion_smoothing = 0.0 with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE): for i in range(SEQ_LENGTH - 1): j = i + 1 depth_i = self.depth[i][:, :, :, 0] depth_j = self.depth[j][:, :, :, 0] image_j = self.image_stack[:, :, :, 3 * j:3 * (j + 1)] image_i = self.image_stack[:, :, :, i * 3:(i + 1) * 3] # We select a pair of consecutive images (and their respective # predicted depth maps). Now we have the network predict a motion # field that connects the two. We feed the pair of images into the # network, once in forward order and then in reverse order. The # results are fed into the loss calculation. The following losses are # calculated: # - RGB and SSIM photometric consistency. # - Cycle consistency of rotations and translations for every pixel. # - L1 smoothness of the disparity and the motion field. # - Depth consistency rot, trans, trans_res, mat = motion_prediction_net.motion_field_net( images=tf.concat([image_i, image_j], axis=-1), weight_reg=self.weight_reg) inv_rot, inv_trans, inv_trans_res, inv_mat = ( motion_prediction_net.motion_field_net( images=tf.concat([image_j, image_i], axis=-1), weight_reg=self.weight_reg)) if self.learn_intrinsics: intrinsic_mat = 0.5 * (mat + inv_mat) else: intrinsic_mat = self.intrinsic_mat[:, 0, :, :] def dilate(x): # Dilation by n pixels is roughtly max pooling by 2 * n + 1. p = self.foreground_dilation * 2 + 1 return tf.nn.max_pool(x, [1, p, p, 1], [1] * 4, 'SAME') trans += trans_res * dilate(self.seg_stack[:, :, :, j:j + 1]) inv_trans += inv_trans_res * dilate( self.seg_stack[:, :, :, i:i + 1]) tf.summary.image('trans%d%d' % (i, i + 1), trans) tf.summary.image('trans%d%d' % (i + 1, i), inv_trans) tf.summary.image('trans_res%d%d' % (i + 1, i), inv_trans_res) tf.summary.image('trans_res%d%d' % (i, i + 1), trans_res) self.motion_smoothing += _smoothness(trans) self.motion_smoothing += _smoothness(inv_trans) tf.summary.scalar( 'trans_stdev', tf.sqrt(0.5 * tf.reduce_mean( tf.square(trans) + tf.square(inv_trans)))) transformed_depth_j = transform_depth_map.using_motion_vector( depth_j, trans, rot, intrinsic_mat) add_result_to_loss_and_summaries( consistency_losses.rgbd_and_motion_consistency_loss( transformed_depth_j, image_j, depth_i, image_i, rot, trans, inv_rot, inv_trans), i, j) transformed_depth_i = transform_depth_map.using_motion_vector( depth_i, inv_trans, inv_rot, intrinsic_mat) add_result_to_loss_and_summaries( consistency_losses.rgbd_and_motion_consistency_loss( transformed_depth_i, image_i, depth_j, image_j, inv_rot, inv_trans, rot, trans), j, i) # Build the total loss as composed of L1 reconstruction, SSIM, smoothing # and object size constraint loss as appropriate. self.reconstr_loss *= self.reconstr_weight self.export('self.reconstr_loss', self.reconstr_loss) self.total_loss = self.reconstr_loss if self.smooth_weight > 0: self.smooth_loss *= self.smooth_weight self.total_loss += self.smooth_loss self.export('self.smooth_loss', self.smooth_loss) if self.ssim_weight > 0: self.ssim_loss *= self.ssim_weight self.total_loss += self.ssim_loss self.export('self.ssim_loss', self.ssim_loss) if self.motion_smoothing_weight > 0: self.motion_smoothing *= self.motion_smoothing_weight self.total_loss += self.motion_smoothing self.export('self.motion_sm_loss', self.motion_smoothing) if self.depth_consistency_loss_weight: self.depth_consistency_loss *= self.depth_consistency_loss_weight self.total_loss += self.depth_consistency_loss self.export('self.depth_consistency_loss', self.depth_consistency_loss) self.rot_loss *= self.rotation_consistency_weight self.trans_loss *= self.translation_consistency_weight self.export('rot_loss', self.rot_loss) self.export('trans_loss', self.trans_loss) self.total_loss += self.rot_loss self.total_loss += self.trans_loss self.export('self.total_loss', self.total_loss)