Merge pull request #2 from soumik12345/zero-dce

Zero-reference Deep Curve Estimation for Low-light Image Enhancement

.gitignore

@@ -131,4 +131,6 @@ dmypy.json
 # Datasets
 datasets/
 **.zip
-**.h5
+**.h5
+lol_dataset_**
+wandb**

Dockerfile

@@ -0,0 +1,17 @@
+# Pull Base Image
+FROM tensorflow/tensorflow:latest-gpu-jupyter
+
+# Set Working Directory
+RUN mkdir /usr/src/enhance-me
+WORKDIR /usr/src/enhance-me
+
+# Set Environment Variables
+ENV PYTHONDONTWRITEBYTECODE 1
+ENV PYTHONUNBUFFERED 1
+
+RUN pip install --upgrade pip setuptools wheel
+RUN pip install gdown matplotlib streamlit tqdm wandb
+
+COPY . /usr/src/enhance-me/
+
+CMD ["jupyter", "notebook", "--port=8888", "--no-browser", "--ip=0.0.0.0", "--allow-root"]

README.md

@@ -8,4 +8,14 @@ app_file: app.py
 pinned: false
 ---
 
-# enhance-me
+# Enhance Me
+
+A unified platform for image enhancement.
+
+## Usage
+
+### Train using Docker
+
+- Build image using `docker build -t enhance-image .`
+
+- Run notebook using `docker run -it --gpus all -p 8888:8888 -v $(pwd):/usr/src/enhance-me enhance-image`

app.py

@@ -1,23 +1,36 @@
+import os
 from PIL import Image
 import streamlit as st
 from tensorflow.keras import utils, backend
 
-from enhance_me.mirnet import MIRNet
+from enhance_me import MIRNet, ZeroDCE
 
 
 def get_mirnet_object() -> MIRNet:
-    mirnet = MIRNet()
-    mirnet.build_model()
     utils.get_file(
         "weights_lol_128.h5",
         "https://github.com/soumik12345/enhance-me/releases/download/v0.2/weights_lol_128.h5",
         cache_dir=".",
         cache_subdir="weights",
     )
+    mirnet = MIRNet()
+    mirnet.build_model()
     mirnet.load_weights("./weights/weights_lol_128.h5")
     return mirnet
 
 
+def get_zero_dce_object(model_alias: str) -> ZeroDCE:
+    utils.get_file(
+        f"{model_alias}.h5",
+        f"https://github.com/soumik12345/enhance-me/releases/download/v0.4/{model_alias}.h5",
+        cache_dir=".",
+        cache_subdir="weights",
+    )
+    dce = ZeroDCE()
+    dce.load_weights(os.path.join("./weights", f"{model_alias}.h5"))
+    return dce
+
+
 def main():
     st.markdown("# Enhance Me")
     st.markdown("Made with :heart: by [geekyRakshit](http://github.com/soumik12345)")
@@ -30,14 +43,39 @@ def main():
             if uploaded_file is not None:
                 original_image = Image.open(uploaded_file)
                 st.image(original_image, caption="original image")
-                st.sidebar.info("Loading MIRNet...")
-                mirnet = get_mirnet_object()
-                st.sidebar.info("Done!")
-                st.sidebar.info("Processing Image...")
-                enhanced_image = mirnet.infer(original_image)
-                st.sidebar.info("Done!")
-                st.image(enhanced_image, caption="enhanced image")
-                backend.clear_session()
+                model_option = st.sidebar.selectbox(
+                    "Please select the model:",
+                    (
+                        "",
+                        "MIRNet",
+                        "Zero-DCE (dce_weights_lol_128)",
+                        "Zero-DCE (dce_weights_lol_128_resize)",
+                        "Zero-DCE (dce_weights_lol_256)",
+                        "Zero-DCE (dce_weights_lol_256_resize)",
+                        "Zero-DCE (dce_weights_unpaired_128)",
+                        "Zero-DCE (dce_weights_unpaired_128_resize)",
+                        "Zero-DCE (dce_weights_unpaired_256)",
+                        "Zero-DCE (dce_weights_unpaired_256_resize)"
+                    ),
+                )
+                if model_option != "":
+                    if model_option == "MIRNet":
+                        st.sidebar.info("Loading MIRNet...")
+                        mirnet = get_mirnet_object()
+                        st.sidebar.info("Done!")
+                        st.sidebar.info("Processing Image...")
+                        enhanced_image = mirnet.infer(original_image)
+                        st.sidebar.info("Done!")
+                        st.image(enhanced_image, caption="enhanced image")
+                    elif "Zero-DCE" in model_option:
+                        model_alias = model_option[model_option.find("(") + 1: model_option.find(")")]
+                        st.sidebar.info("Loading Zero-DCE...")
+                        zero_dce = get_zero_dce_object(model_alias)
+                        st.sidebar.info("Done!")
+                        enhanced_image = zero_dce.infer(original_image)
+                        st.sidebar.info("Done!")
+                        st.image(enhanced_image, caption="enhanced image")
+                    backend.clear_session()
 
 
 if __name__ == "__main__":

enhance_me/__init__.py

@@ -0,0 +1,2 @@
+from .mirnet import MIRNet
+from .zero_dce import ZeroDCE

enhance_me/augmentation.py

@@ -49,3 +49,38 @@ class AugmentationFactory:
         return tf.image.rot90(input_image, condition), tf.image.rot90(
             enhanced_image, condition
         )
+
+
+class UnpairedAugmentationFactory:
+    def __init__(self, image_size) -> None:
+        self.image_size = image_size
+
+    def random_crop(self, image):
+        image_shape = tf.shape(image)[:2]
+        crop_w = tf.random.uniform(
+            shape=(), maxval=image_shape[1] - self.image_size + 1, dtype=tf.int32
+        )
+        crop_h = tf.random.uniform(
+            shape=(), maxval=image_shape[0] - self.image_size + 1, dtype=tf.int32
+        )
+        return image[
+            crop_h : crop_h + self.image_size, crop_w : crop_w + self.image_size
+        ]
+
+    def random_horizontal_flip(self, image):
+        return tf.cond(
+            tf.random.uniform(shape=(), maxval=1) < 0.5,
+            lambda: image,
+            lambda: tf.image.flip_left_right(image),
+        )
+
+    def random_vertical_flip(self, image):
+        return tf.cond(
+            tf.random.uniform(shape=(), maxval=1) < 0.5,
+            lambda: image,
+            lambda: tf.image.flip_up_down(image),
+        )
+
+    def random_rotate(self, image):
+        condition = tf.random.uniform(shape=(), maxval=4, dtype=tf.int32)
+        return tf.image.rot90(image, condition)

enhance_me/commons.py

@@ -61,3 +61,18 @@ def download_lol_dataset():
     test_enhanced_images = sorted(glob("./datasets/lol_dataset/eval15/high/*"))
     assert len(test_low_images) == len(test_enhanced_images)
     return (low_images, enhanced_images), (test_low_images, test_enhanced_images)
+
+
+def download_unpaired_low_light_dataset():
+    utils.get_file(
+        "low_light_dataset.zip",
+        "https://github.com/soumik12345/enhance-me/releases/download/v0.3/low_light_dataset.zip",
+        cache_dir="./",
+        cache_subdir="./datasets",
+        extract=True,
+    )
+    low_images = glob("./datasets/low_light_dataset/*.png")
+    test_low_images = sorted(glob("./datasets/low_light_dataset/eval15/low/*"))
+    test_enhanced_images = sorted(glob("./datasets/low_light_dataset/eval15/high/*"))
+    assert len(test_low_images) == len(test_enhanced_images)
+    return low_images, (test_low_images, test_enhanced_images)

enhance_me/zero_dce/__init__.py

@@ -0,0 +1 @@
+from .zero_dce import ZeroDCE

enhance_me/zero_dce/dataloader.py

@@ -0,0 +1,79 @@
+import tensorflow as tf
+from typing import List
+
+from ..commons import read_image
+from ..augmentation import UnpairedAugmentationFactory
+
+
+class UnpairedLowLightDataset:
+    def __init__(
+        self,
+        image_size: int = 256,
+        apply_resize: bool = False,
+        apply_random_horizontal_flip: bool = True,
+        apply_random_vertical_flip: bool = True,
+        apply_random_rotation: bool = True,
+    ) -> None:
+        self.augmentation_factory = UnpairedAugmentationFactory(image_size=image_size)
+        self.image_size = image_size
+        self.apply_resize = apply_resize
+        self.apply_random_horizontal_flip = apply_random_horizontal_flip
+        self.apply_random_vertical_flip = apply_random_vertical_flip
+        self.apply_random_rotation = apply_random_rotation
+
+    def _resize(self, image):
+        return tf.image.resize(image, (self.image_size, self.image_size))
+
+    def _get_dataset(self, images: List[str], batch_size: int, is_train: bool):
+        dataset = tf.data.Dataset.from_tensor_slices((images))
+        dataset = dataset.map(read_image, num_parallel_calls=tf.data.AUTOTUNE)
+        dataset = (
+            dataset.map(
+                self.augmentation_factory.random_crop,
+                num_parallel_calls=tf.data.AUTOTUNE,
+            )
+            if not self.apply_resize
+            else dataset.map(self._resize, num_parallel_calls=tf.data.AUTOTUNE)
+        )
+        if is_train:
+            dataset = (
+                dataset.map(
+                    self.augmentation_factory.random_horizontal_flip,
+                    num_parallel_calls=tf.data.AUTOTUNE,
+                )
+                if self.apply_random_horizontal_flip
+                else dataset
+            )
+            dataset = (
+                dataset.map(
+                    self.augmentation_factory.random_vertical_flip,
+                    num_parallel_calls=tf.data.AUTOTUNE,
+                )
+                if self.apply_random_vertical_flip
+                else dataset
+            )
+            dataset = (
+                dataset.map(
+                    self.augmentation_factory.random_rotate,
+                    num_parallel_calls=tf.data.AUTOTUNE,
+                )
+                if self.apply_random_rotation
+                else dataset
+            )
+        dataset = dataset.batch(batch_size, drop_remainder=True)
+        return dataset
+
+    def get_datasets(
+        self,
+        images: List[str],
+        val_split: float = 0.2,
+        batch_size: int = 16,
+    ):
+        split_index = int(len(images) * (1 - val_split))
+        train_images = images[:split_index]
+        val_images = images[split_index:]
+        print(f"Number of train data points: {len(train_images)}")
+        print(f"Number of validation data points: {len(val_images)}")
+        train_dataset = self._get_dataset(train_images, batch_size, is_train=True)
+        val_dataset = self._get_dataset(val_images, batch_size, is_train=False)
+        return train_dataset, val_dataset

enhance_me/zero_dce/dce_net.py

@@ -0,0 +1,31 @@
+import tensorflow as tf
+from tensorflow.keras import layers, Input, Model
+
+
+def build_dce_net() -> Model:
+    input_image = Input(shape=[None, None, 3])
+    conv1 = layers.Conv2D(
+        32, (3, 3), strides=(1, 1), activation="relu", padding="same"
+    )(input_image)
+    conv2 = layers.Conv2D(
+        32, (3, 3), strides=(1, 1), activation="relu", padding="same"
+    )(conv1)
+    conv3 = layers.Conv2D(
+        32, (3, 3), strides=(1, 1), activation="relu", padding="same"
+    )(conv2)
+    conv4 = layers.Conv2D(
+        32, (3, 3), strides=(1, 1), activation="relu", padding="same"
+    )(conv3)
+    int_con1 = layers.Concatenate(axis=-1)([conv4, conv3])
+    conv5 = layers.Conv2D(
+        32, (3, 3), strides=(1, 1), activation="relu", padding="same"
+    )(int_con1)
+    int_con2 = layers.Concatenate(axis=-1)([conv5, conv2])
+    conv6 = layers.Conv2D(
+        32, (3, 3), strides=(1, 1), activation="relu", padding="same"
+    )(int_con2)
+    int_con3 = layers.Concatenate(axis=-1)([conv6, conv1])
+    x_r = layers.Conv2D(24, (3, 3), strides=(1, 1), activation="tanh", padding="same")(
+        int_con3
+    )
+    return Model(inputs=input_image, outputs=x_r)

enhance_me/zero_dce/losses/__init__.py

@@ -0,0 +1,36 @@
+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

enhance_me/zero_dce/zero_dce.py

@@ -0,0 +1,183 @@
+import os
+import numpy as np
+from PIL import Image
+from datetime import datetime
+
+import tensorflow as tf
+from tensorflow import keras
+from tensorflow.keras import optimizers, mixed_precision, Model
+from wandb.keras import WandbCallback
+
+from .dce_net import build_dce_net
+from .dataloader import UnpairedLowLightDataset
+from .losses import (
+    color_constancy_loss,
+    exposure_loss,
+    illumination_smoothness_loss,
+    SpatialConsistencyLoss,
+)
+from ..commons import (
+    download_lol_dataset,
+    download_unpaired_low_light_dataset,
+    init_wandb,
+)
+
+
+class ZeroDCE(Model):
+    def __init__(
+        self,
+        experiment_name=None,
+        wandb_api_key=None,
+        use_mixed_precision: bool = False,
+        **kwargs
+    ):
+        super(ZeroDCE, self).__init__(**kwargs)
+        self.experiment_name = experiment_name
+        if use_mixed_precision:
+            policy = mixed_precision.Policy("mixed_float16")
+            mixed_precision.set_global_policy(policy)
+        if wandb_api_key is not None:
+            init_wandb("zero-dce", experiment_name, wandb_api_key)
+            self.using_wandb = True
+        else:
+            self.using_wandb = False
+        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,
+        )
+
+    def build_datasets(
+        self,
+        image_size: int = 256,
+        dataset_label: str = "lol",
+        apply_resize: bool = False,
+        apply_random_horizontal_flip: bool = True,
+        apply_random_vertical_flip: bool = True,
+        apply_random_rotation: bool = True,
+        val_split: float = 0.2,
+        batch_size: int = 16,
+    ) -> None:
+        if dataset_label == "lol":
+            (self.low_images, _), (self.test_low_images, _) = download_lol_dataset()
+        elif dataset_label == "unpaired":
+            self.low_images, (
+                self.test_low_images,
+                _,
+            ) = download_unpaired_low_light_dataset()
+        data_loader = UnpairedLowLightDataset(
+            image_size,
+            apply_resize,
+            apply_random_horizontal_flip,
+            apply_random_vertical_flip,
+            apply_random_rotation,
+        )
+        self.train_dataset, self.val_dataset = data_loader.get_datasets(
+            self.low_images, val_split, batch_size
+        )
+
+    def train(self, epochs: int):
+        log_dir = os.path.join(
+            self.experiment_name,
+            "logs",
+            datetime.now().strftime("%Y%m%d-%H%M%S"),
+        )
+        tensorboard_callback = keras.callbacks.TensorBoard(log_dir, histogram_freq=1)
+        callbacks = [tensorboard_callback]
+        if self.using_wandb:
+            callbacks += [WandbCallback()]
+        history = self.fit(
+            self.train_dataset,
+            validation_data=self.val_dataset,
+            epochs=epochs,
+            callbacks=callbacks,
+        )
+        return history
+
+    def infer(self, original_image):
+        image = keras.preprocessing.image.img_to_array(original_image)
+        image = image.astype("float32") / 255.0
+        image = np.expand_dims(image, axis=0)
+        output_image = self.call(image)
+        output_image = tf.cast((output_image[0, :, :, :] * 255), dtype=np.uint8)
+        output_image = Image.fromarray(output_image.numpy())
+        return output_image
+
+    def infer_from_file(self, original_image_file: str):
+        original_image = Image.open(original_image_file)
+        return self.infer(original_image)

notebooks/enhance_me_train.ipynb

@@ -37,11 +37,12 @@
     "import os\n",
     "import sys\n",
     "\n",
-    "sys.path.append(\"./enhance-me\")\n",
+    "sys.path.append(\"..\")\n",
     "\n",
     "from PIL import Image\n",
     "from enhance_me import commons\n",
-    "from enhance_me.mirnet import MIRNet"
+    "from enhance_me.mirnet import MIRNet\n",
+    "from enhance_me.zero_dce import ZeroDCE"
    ]
   },
   {
@@ -170,7 +171,7 @@
     "    enhanced_image = mirnet.infer(original_image)\n",
     "    ground_truth = Image.open(mirnet.test_enhanced_images[index])\n",
     "    commons.plot_results(\n",
-    "        [original_image, ground_truth, ground_truth],\n",
+    "        [original_image, ground_truth, enhanced_image],\n",
     "        [\"Original Image\", \"Ground Truth\", \"Enhanced Image\"],\n",
     "        (18, 18),\n",
     "    )"
@@ -183,6 +184,92 @@
     "id": "dO-IbNQHkB3R"
    },
    "outputs": [],
+   "source": [
+    "# @title Zero-DCE Train Configs\n",
+    "\n",
+    "experiment_name = \"unpaired_low_light_256_resize\"  # @param {type:\"string\"}\n",
+    "image_size = 256  # @param {type:\"integer\"}\n",
+    "dataset_label = \"unpaired\"  # @param [\"lol\", \"unpaired\"]\n",
+    "use_mixed_precision = False  # @param {type:\"boolean\"}\n",
+    "apply_resize = True  # @param {type:\"boolean\"}\n",
+    "apply_random_horizontal_flip = True  # @param {type:\"boolean\"}\n",
+    "apply_random_vertical_flip = True  # @param {type:\"boolean\"}\n",
+    "apply_random_rotation = True  # @param {type:\"boolean\"}\n",
+    "wandb_api_key = \"\"  # @param {type:\"string\"}\n",
+    "val_split = 0.1  # @param {type:\"slider\", min:0.1, max:1.0, step:0.1}\n",
+    "batch_size = 16  # @param {type:\"integer\"}\n",
+    "learning_rate = 1e-4  # @param {type:\"number\"}\n",
+    "epsilon = 1e-3  # @param {type:\"number\"}\n",
+    "epochs = 100  # @param {type:\"slider\", min:10, max:100, step:5}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "zero_dce = ZeroDCE(\n",
+    "    experiment_name=experiment_name,\n",
+    "    wandb_api_key=None if wandb_api_key == \"\" else wandb_api_key,\n",
+    "    use_mixed_precision=use_mixed_precision\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "zero_dce.build_datasets(\n",
+    "    image_size=image_size,\n",
+    "    dataset_label=dataset_label,\n",
+    "    apply_resize=apply_resize,\n",
+    "    apply_random_horizontal_flip=apply_random_horizontal_flip,\n",
+    "    apply_random_vertical_flip=apply_random_vertical_flip,\n",
+    "    apply_random_rotation=apply_random_rotation,\n",
+    "    val_split=val_split,\n",
+    "    batch_size=batch_size\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "scrolled": false
+   },
+   "outputs": [],
+   "source": [
+    "zero_dce.compile(learning_rate=learning_rate)\n",
+    "history = zero_dce.train(epochs=epochs)\n",
+    "zero_dce.save_weights(os.path.join(experiment_name, \"weights.h5\"))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "scrolled": false
+   },
+   "outputs": [],
+   "source": [
+    "for index, low_image_file in enumerate(zero_dce.test_low_images):\n",
+    "    original_image = Image.open(low_image_file)\n",
+    "    enhanced_image = zero_dce.infer(original_image)\n",
+    "    commons.plot_results(\n",
+    "        [original_image, enhanced_image],\n",
+    "        [\"Original Image\", \"Enhanced Image\"],\n",
+    "        (18, 18),\n",
+    "    )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
    "source": []
   }
  ],
@@ -197,13 +284,23 @@
    "provenance": []
   },
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
    "name": "python3"
   },
   "language_info": {
-   "name": "python"
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.10"
   }
  },
  "nbformat": 4,
- "nbformat_minor": 0
+ "nbformat_minor": 1
 }

test.py

@@ -0,0 +1,4 @@
+from enhance_me.commons import download_unpaired_low_light_dataset
+
+
+download_unpaired_low_light_dataset()