added loss functions

enhance_me/zero_dce/losses/__init__.py

@@ -0,0 +1,32 @@
+import tensorflow as tf
+
+from .spatial_constancy import SpatialConsistencyLoss
+
+
+def color_constancy_loss(x):
+    mean_rgb = tf.reduce_mean(x, axis=(1, 2), keepdims=True)
+    mean_r, mean_g, mean_b = mean_rgb[:, :, :, 0], mean_rgb[:, :, :, 1], mean_rgb[:, :, :, 2]
+    diff_rg = tf.square(mean_r - mean_g)
+    diff_rb = tf.square(mean_r - mean_b)
+    diff_gb = tf.square(mean_b - mean_g)
+    return tf.sqrt(tf.square(diff_rg) + tf.square(diff_rb) + tf.square(diff_gb))
+
+
+def exposure_loss(x, mean_val=0.6):
+    x = tf.reduce_mean(x, axis=3, keepdims=True)
+    mean = tf.nn.avg_pool2d(x, ksize=16, strides=16, padding="VALID")
+    return tf.reduce_mean(tf.square(mean - mean_val))
+
+
+def illumination_smoothness_loss(x):
+    batch_size = tf.shape(x)[0]
+    h_x = tf.shape(x)[1]
+    w_x = tf.shape(x)[2]
+    count_h = (tf.shape(x)[2] - 1) * tf.shape(x)[3]
+    count_w = tf.shape(x)[2] * (tf.shape(x)[3] - 1)
+    h_tv = tf.reduce_sum(tf.square((x[:, 1:, :, :] - x[:, : h_x - 1, :, :])))
+    w_tv = tf.reduce_sum(tf.square((x[:, :, 1:, :] - x[:, :, : w_x - 1, :])))
+    batch_size = tf.cast(batch_size, dtype=tf.float32)
+    count_h = tf.cast(count_h, dtype=tf.float32)
+    count_w = tf.cast(count_w, dtype=tf.float32)
+    return 2 * (h_tv / count_h + w_tv / count_w) / batch_size

enhance_me/zero_dce/losses/spatial_constancy.py

@@ -0,0 +1,63 @@
+import tensorflow as tf
+from tensorflow.keras import losses
+
+
+class SpatialConsistencyLoss(losses.Loss):
+    def __init__(self, **kwargs):
+        super(SpatialConsistencyLoss, self).__init__(reduction="none")
+
+        self.left_kernel = tf.constant(
+            [[[[0, 0, 0]], [[-1, 1, 0]], [[0, 0, 0]]]], dtype=tf.float32
+        )
+        self.right_kernel = tf.constant(
+            [[[[0, 0, 0]], [[0, 1, -1]], [[0, 0, 0]]]], dtype=tf.float32
+        )
+        self.up_kernel = tf.constant(
+            [[[[0, -1, 0]], [[0, 1, 0]], [[0, 0, 0]]]], dtype=tf.float32
+        )
+        self.down_kernel = tf.constant(
+            [[[[0, 0, 0]], [[0, 1, 0]], [[0, -1, 0]]]], dtype=tf.float32
+        )
+
+    def call(self, y_true, y_pred):
+
+        original_mean = tf.reduce_mean(y_true, 3, keepdims=True)
+        enhanced_mean = tf.reduce_mean(y_pred, 3, keepdims=True)
+        original_pool = tf.nn.avg_pool2d(
+            original_mean, ksize=4, strides=4, padding="VALID"
+        )
+        enhanced_pool = tf.nn.avg_pool2d(
+            enhanced_mean, ksize=4, strides=4, padding="VALID"
+        )
+
+        d_original_left = tf.nn.conv2d(
+            original_pool, self.left_kernel, strides=[1, 1, 1, 1], padding="SAME"
+        )
+        d_original_right = tf.nn.conv2d(
+            original_pool, self.right_kernel, strides=[1, 1, 1, 1], padding="SAME"
+        )
+        d_original_up = tf.nn.conv2d(
+            original_pool, self.up_kernel, strides=[1, 1, 1, 1], padding="SAME"
+        )
+        d_original_down = tf.nn.conv2d(
+            original_pool, self.down_kernel, strides=[1, 1, 1, 1], padding="SAME"
+        )
+
+        d_enhanced_left = tf.nn.conv2d(
+            enhanced_pool, self.left_kernel, strides=[1, 1, 1, 1], padding="SAME"
+        )
+        d_enhanced_right = tf.nn.conv2d(
+            enhanced_pool, self.right_kernel, strides=[1, 1, 1, 1], padding="SAME"
+        )
+        d_enhanced_up = tf.nn.conv2d(
+            enhanced_pool, self.up_kernel, strides=[1, 1, 1, 1], padding="SAME"
+        )
+        d_enhanced_down = tf.nn.conv2d(
+            enhanced_pool, self.down_kernel, strides=[1, 1, 1, 1], padding="SAME"
+        )
+
+        d_left = tf.square(d_original_left - d_enhanced_left)
+        d_right = tf.square(d_original_right - d_enhanced_right)
+        d_up = tf.square(d_original_up - d_enhanced_up)
+        d_down = tf.square(d_original_down - d_enhanced_down)
+        return d_left + d_right + d_up + d_down