added zero-dce model

enhance_me/zero_dce/dataloader.py

@@ -17,7 +17,7 @@ class UnpairedLowLightDataset:
         self.apply_random_vertical_flip = apply_random_vertical_flip
         self.apply_random_rotation = apply_random_rotation
 
-    def load_data(self, image_path):
+    def _load_data(self, image_path):
         image = tf.io.read_file(image_path)
         image = tf.image.decode_png(image, channels=3)
         image = image / 255.0
@@ -25,7 +25,7 @@ class UnpairedLowLightDataset:
 
     def _get_dataset(self, images: List[str], batch_size: int, is_train: bool):
         dataset = tf.data.Dataset.from_tensor_slices((images))
-        dataset = dataset.map(self.load_data, num_parallel_calls=tf.data.AUTOTUNE)
+        dataset = dataset.map(self._load_data, num_parallel_calls=tf.data.AUTOTUNE)
         dataset = dataset.map(
             self.augmentation_factory.random_crop, num_parallel_calls=tf.data.AUTOTUNE
         )

enhance_me/zero_dce/losses/__init__.py

@@ -5,7 +5,11 @@ 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]
+    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)

enhance_me/zero_dce/models/zero_dce.py

@@ -0,0 +1,96 @@
+import tensorflow as tf
+from tensorflow.keras import optimizers, Model
+
+from .dce_net import build_dce_net
+from ..dataloader import UnpairedLowLightDataset
+from ..losses import (
+    color_constancy_loss,
+    exposure_loss,
+    illumination_smoothness_loss,
+    SpatialConsistencyLoss,
+)
+
+
+class ZeroDCE(Model):
+    def __init__(self, **kwargs):
+        super(ZeroDCE, self).__init__(**kwargs)
+        self.dce_model = build_dce_net()
+
+    def compile(self, learning_rate, **kwargs):
+        super(ZeroDCE, self).compile(**kwargs)
+        self.optimizer = optimizers.Adam(learning_rate=learning_rate)
+        self.spatial_constancy_loss = SpatialConsistencyLoss(reduction="none")
+
+    def get_enhanced_image(self, data, output):
+        r1 = output[:, :, :, :3]
+        r2 = output[:, :, :, 3:6]
+        r3 = output[:, :, :, 6:9]
+        r4 = output[:, :, :, 9:12]
+        r5 = output[:, :, :, 12:15]
+        r6 = output[:, :, :, 15:18]
+        r7 = output[:, :, :, 18:21]
+        r8 = output[:, :, :, 21:24]
+        x = data + r1 * (tf.square(data) - data)
+        x = x + r2 * (tf.square(x) - x)
+        x = x + r3 * (tf.square(x) - x)
+        enhanced_image = x + r4 * (tf.square(x) - x)
+        x = enhanced_image + r5 * (tf.square(enhanced_image) - enhanced_image)
+        x = x + r6 * (tf.square(x) - x)
+        x = x + r7 * (tf.square(x) - x)
+        enhanced_image = x + r8 * (tf.square(x) - x)
+        return enhanced_image
+
+    def call(self, data):
+        dce_net_output = self.dce_model(data)
+        return self.get_enhanced_image(data, dce_net_output)
+
+    def compute_losses(self, data, output):
+        enhanced_image = self.get_enhanced_image(data, output)
+        loss_illumination = 200 * illumination_smoothness_loss(output)
+        loss_spatial_constancy = tf.reduce_mean(
+            self.spatial_constancy_loss(enhanced_image, data)
+        )
+        loss_color_constancy = 5 * tf.reduce_mean(color_constancy_loss(enhanced_image))
+        loss_exposure = 10 * tf.reduce_mean(exposure_loss(enhanced_image))
+        total_loss = (
+            loss_illumination
+            + loss_spatial_constancy
+            + loss_color_constancy
+            + loss_exposure
+        )
+        return {
+            "total_loss": total_loss,
+            "illumination_smoothness_loss": loss_illumination,
+            "spatial_constancy_loss": loss_spatial_constancy,
+            "color_constancy_loss": loss_color_constancy,
+            "exposure_loss": loss_exposure,
+        }
+
+    def train_step(self, data):
+        with tf.GradientTape() as tape:
+            output = self.dce_model(data)
+            losses = self.compute_losses(data, output)
+        gradients = tape.gradient(
+            losses["total_loss"], self.dce_model.trainable_weights
+        )
+        self.optimizer.apply_gradients(zip(gradients, self.dce_model.trainable_weights))
+        return losses
+
+    def test_step(self, data):
+        output = self.dce_model(data)
+        return self.compute_losses(data, output)
+
+    def save_weights(self, filepath, overwrite=True, save_format=None, options=None):
+        """While saving the weights, we simply save the weights of the DCE-Net"""
+        self.dce_model.save_weights(
+            filepath, overwrite=overwrite, save_format=save_format, options=options
+        )
+
+    def load_weights(self, filepath, by_name=False, skip_mismatch=False, options=None):
+        """While loading the weights, we simply load the weights of the DCE-Net"""
+        self.dce_model.load_weights(
+            filepath=filepath,
+            by_name=by_name,
+            skip_mismatch=skip_mismatch,
+            options=options,
+        )