From 26dda685232f9f1e6ebe3ccf29b5a2436300e694 Mon Sep 17 00:00:00 2001 From: Michael Pivato Date: Thu, 5 Aug 2021 08:20:31 +0000 Subject: [PATCH] Add Smooth Loss --- unsupervised/loss.py | 31 +++++++++++++++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/unsupervised/loss.py b/unsupervised/loss.py index d62e5b1..c4057a1 100644 --- a/unsupervised/loss.py +++ b/unsupervised/loss.py @@ -51,3 +51,34 @@ def make_combined_ssim_l1_loss(ssim_weight: int = 0.85, other_loss_fn=l1_loss): return ssim_weight * ssim + (1 - ssim_weight) * other_loss_fn(target_img, reprojected_img) return combined_ssim_loss + + +# TODO: Consider other gradient methods for calculating smoothness, e.g. convolution methods such as Sobel +def smooth_loss(depth, colour_image): + """ + Calculate the edge-aware per-pixel smooth loss on a depth map, with image scaled appropriately to the depth map + + Does this equation (equation 3 in monodepth2 paper): + |dxd*t|e^(-|dxIt|) + |dyd*t|e^(-|dyIt|) + + :param depth: Tensor with shape (B, h, w, 1) - disparity, such as the depth map + :param colour_image: Tensor with shape (B, h, w, 3) - colour image, same resolution as disparity map + :return: smooth loss + """ + # Mean normalised inverse depth + normalised_depth = depth / (tf.reduce_mean(depth, [1, 2], keepdims=True) + 1e-7) + + # Nothing fancy here for gradients (follows sfmlearner/monodepth), just shift 1 pixel and + # compare the change (x/y shift left/up 1 pixel) + depth_gradient_x = tf.abs(normalised_depth[:, :-1, :, :] - normalised_depth[:, 1:, :, :]) + depth_gradient_y = tf.abs(normalised_depth[:, :, :-1, :] - normalised_depth[:, :, 1:, :]) + + # Colour gradients to work better with edges, monodepth 1/2 uses these + image_gradient_x = tf.abs(colour_image[:, :-1, :, :] - colour_image[:, 1:, :, :]) + image_gradient_y = tf.abs(colour_image[:, :, :-1, :] - colour_image[:, :, 1:, :]) + + # Average the 3 colour channels into a single channel, so can be compared with the depth disparities + smooth_x = depth_gradient_x * tf.exp(-tf.reduce_mean(image_gradient_x, 3, keepdims=True)) + smooth_y = depth_gradient_y * tf.exp(-tf.reduce_mean(image_gradient_y, 3, keepdims=True)) + + return tf.reduce_mean(smooth_x) + tf.reduce_mean(smooth_y)