title zerodce for lowlight image enhancement soumik rakshithttp github comsoumik12345 converted to keras 3 by soumik rakshithttp github comsoumik12345 date created 20210918 last modified 20230715 description implementing zeroreference deep curve estimation for lowlight image enhancement accelerator gpu introduction zeroreference deep curve estimation or zerodce formulates lowlight image enhancement as the task of estimating an imagespecific tonal curvehttps en wikipedia orgwikicurvetonality with a deep neural network in this example we train a lightweight deep network dcenet to estimate pixelwise and highorder tonal curves for dynamic range adjustment of a given image zerodce takes a lowlight image as input and produces highorder tonal curves as its output these curves are then used for pixelwise adjustment on the dynamic range of the input to obtain an enhanced image the curve estimation process is done in such a way that it maintains the range of the enhanced image and preserves the contrast of neighboring pixels this curve estimation is inspired by curves adjustment used in photo editing software such as adobe photoshop where users can adjust points throughout an images tonal range zerodce is appealing because of its relaxed assumptions with regard to reference images it does not require any inputoutput image pairs during training this is achieved through a set of carefully formulated nonreference loss functions which implicitly measure the enhancement quality and guide the training of the network references zeroreference deep curve estimation for lowlight image enhancementhttps arxiv orgpdf2001 06826 pdf curves adjustment in adobe photoshophttps helpx adobe comphotoshopusingcurvesadjustment html downloading loldataset the lol dataset has been created for lowlight image enhancement it provides 485 images for training and 15 for testing each image pair in the dataset consists of a lowlight input image and its corresponding wellexposed reference image shell wget https huggingface codatasetsgeekyrakshitloldatasetresolvemainloldataset zip unzip q loldataset zip rm loldataset zip creating a tensorflow dataset we use 300 lowlight images from the lol dataset training set for training and we use the remaining 185 lowlight images for validation we resize the images to size 256 x 256 to be used for both training and validation note that in order to train the dcenet we will not require the corresponding enhanced images the zerodce framework the goal of dcenet is to estimate a set of bestfitting lightenhancement curves lecurves given an input image the framework then maps all pixels of the inputs rgb channels by applying the curves iteratively to obtain the final enhanced image understanding lightenhancement curves a lighenhancement curve is a kind of curve that can map a lowlight image to its enhanced version automatically where the selfadaptive curve parameters are solely dependent on the input image when designing such a curve three objectives should be taken into account each pixel value of the enhanced image should be in the normalized range 0 1 in order to avoid information loss induced by overflow truncation it should be monotonous to preserve the contrast between neighboring pixels the shape of this curve should be as simple as possible and the curve should be differentiable to allow backpropagation the lightenhancement curve is separately applied to three rgb channels instead of solely on the illumination channel the threechannel adjustment can better preserve the inherent color and reduce the risk of oversaturation https lichongyi github iozerodcefilesframework png dcenet the dcenet is a lightweight deep neural network that learns the mapping between an input image and its bestfitting curve parameter maps the input to the dcenet is a lowlight image while the outputs are a set of pixelwise curve parameter maps for corresponding higherorder curves it is a plain cnn of seven convolutional layers with symmetrical concatenation each layer consists of 32 convolutional kernels of size 33 and stride 1 followed by the relu activation function the last convolutional layer is followed by the tanh activation function which produces 24 parameter maps for 8 iterations where each iteration requires three curve parameter maps for the three channels https i imgur comhtig34w png loss functions to enable zeroreference learning in dcenet we use a set of differentiable zeroreference losses that allow us to evaluate the quality of enhanced images color constancy loss the color constancy loss is used to correct the potential color deviations in the enhanced image exposure loss to restrain underoverexposed regions we use the exposure control loss it measures the distance between the average intensity value of a local region and a preset wellexposedness level set to 0 6 illumination smoothness loss to preserve the monotonicity relations between neighboring pixels the illumination smoothness loss is added to each curve parameter map spatial consistency loss the spatial consistency loss encourages spatial coherence of the enhanced image by preserving the contrast between neighboring regions across the input image and its enhanced version deep curve estimation model we implement the zerodce framework as a keras subclassed model while saving the weights we simply save the weights of the dcenet self dcemodel saveweights filepath overwriteoverwrite saveformatsaveformat optionsoptions def loadweights self filepath bynamefalse skipmismatchfalse optionsnone training inference inference on test images we compare the test images from loldataset enhanced by mirnet with images enhanced via the pil imageops autocontrast function you can use the trained model hosted on hugging face hubhttps huggingface cokerasiolowlightimageenhancement and try the demo on hugging face spaceshttps huggingface cospaceskerasiolowlightimageenhancement introduction zero reference deep curve estimation or zero dce formulates low light image enhancement as the task of estimating an image specific tonal curve https en wikipedia org wiki curve_ tonality with a deep neural network in this example we train a lightweight deep network dce net to estimate pixel wise and high order tonal curves for dynamic range adjustment of a given image zero dce takes a low light image as input and produces high order tonal curves as its output these curves are then used for pixel wise adjustment on the dynamic range of the input to obtain an enhanced image the curve estimation process is done in such a way that it maintains the range of the enhanced image and preserves the contrast of neighboring pixels this curve estimation is inspired by curves adjustment used in photo editing software such as adobe photoshop where users can adjust points throughout an image s tonal range zero dce is appealing because of its relaxed assumptions with regard to reference images it does not require any input output image pairs during training this is achieved through a set of carefully formulated non reference loss functions which implicitly measure the enhancement quality and guide the training of the network references zero reference deep curve estimation for low light image enhancement https arxiv org pdf 2001 06826 pdf curves adjustment in adobe photoshop https helpx adobe com photoshop using curves adjustment html downloading loldataset the lol dataset has been created for low light image enhancement it provides 485 images for training and 15 for testing each image pair in the dataset consists of a low light input image and its corresponding well exposed reference image shell wget https huggingface co datasets geekyrakshit lol dataset resolve main lol_dataset zip unzip q lol_dataset zip rm lol_dataset zip creating a tensorflow dataset we use 300 low light images from the lol dataset training set for training and we use the remaining 185 low light images for validation we resize the images to size 256 x 256 to be used for both training and validation note that in order to train the dce net we will not require the corresponding enhanced images the zero dce framework the goal of dce net is to estimate a set of best fitting light enhancement curves le curves given an input image the framework then maps all pixels of the input s rgb channels by applying the curves iteratively to obtain the final enhanced image understanding light enhancement curves a ligh enhancement curve is a kind of curve that can map a low light image to its enhanced version automatically where the self adaptive curve parameters are solely dependent on the input image when designing such a curve three objectives should be taken into account each pixel value of the enhanced image should be in the normalized range 0 1 in order to avoid information loss induced by overflow truncation it should be monotonous to preserve the contrast between neighboring pixels the shape of this curve should be as simple as possible and the curve should be differentiable to allow backpropagation the light enhancement curve is separately applied to three rgb channels instead of solely on the illumination channel the three channel adjustment can better preserve the inherent color and reduce the risk of over saturation https li chongyi github io zero dce_files framework png dce net the dce net is a lightweight deep neural network that learns the mapping between an input image and its best fitting curve parameter maps the input to the dce net is a low light image while the outputs are a set of pixel wise curve parameter maps for corresponding higher order curves it is a plain cnn of seven convolutional layers with symmetrical concatenation each layer consists of 32 convolutional kernels of size 3 3 and stride 1 followed by the relu activation function the last convolutional layer is followed by the tanh activation function which produces 24 parameter maps for 8 iterations where each iteration requires three curve parameter maps for the three channels https i imgur com htig34w png loss functions to enable zero reference learning in dce net we use a set of differentiable zero reference losses that allow us to evaluate the quality of enhanced images color constancy loss the color constancy loss is used to correct the potential color deviations in the enhanced image exposure loss to restrain under over exposed regions we use the exposure control loss it measures the distance between the average intensity value of a local region and a preset well exposedness level set to 0 6 illumination smoothness loss to preserve the monotonicity relations between neighboring pixels the illumination smoothness loss is added to each curve parameter map spatial consistency loss the spatial consistency loss encourages spatial coherence of the enhanced image by preserving the contrast between neighboring regions across the input image and its enhanced version deep curve estimation model we implement the zero dce framework as a keras subclassed model while saving the weights we simply save the weights of the dce net while loading the weights we simply load the weights of the dce net training inference inference on test images we compare the test images from loldataset enhanced by mirnet with images enhanced via the pil imageops autocontrast function you can use the trained model hosted on hugging face hub https huggingface co keras io low light image enhancement and try the demo on hugging face spaces https huggingface co spaces keras io low light image enhancement
import os
import random
import numpy as np
from glob import glob
from PIL import Image, ImageOps
import matplotlib.pyplot as plt

import keras
from keras import layers

import tensorflow as tf





IMAGE_SIZE = 256
BATCH_SIZE = 16
MAX_TRAIN_IMAGES = 400


def load_data(image_path):
    image = tf.io.read_file(image_path)
    image = tf.image.decode_png(image, channels=3)
    image = tf.image.resize(images=image, size=[IMAGE_SIZE, IMAGE_SIZE])
    image = image / 255.0
    return image


def data_generator(low_light_images):
    dataset = tf.data.Dataset.from_tensor_slices((low_light_images))
    dataset = dataset.map(load_data, num_parallel_calls=tf.data.AUTOTUNE)
    dataset = dataset.batch(BATCH_SIZE, drop_remainder=True)
    return dataset


train_low_light_images = sorted(glob("./lol_dataset/our485/low/*"))[
    :MAX_TRAIN_IMAGES
]
val_low_light_images = sorted(glob("./lol_dataset/our485/low/*"))[
    MAX_TRAIN_IMAGES:
]
test_low_light_images = sorted(glob("./lol_dataset/eval15/low/*"))


train_dataset = data_generator(train_low_light_images)
val_dataset = data_generator(val_low_light_images)

print("Train Dataset:", train_dataset)
print("Validation Dataset:", val_dataset)




def build_dce_net():
    input_img = keras.Input(shape=[None, None, 3])
    conv1 = layers.Conv2D(
        32, (3, 3), strides=(1, 1), activation="relu", padding="same"
    )(input_img)
    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 keras.Model(inputs=input_img, outputs=x_r)







def color_constancy_loss(x):
    mean_rgb = tf.reduce_mean(x, axis=(1, 2), keepdims=True)
    mr, mg, mb = (
        mean_rgb[:, :, :, 0],
        mean_rgb[:, :, :, 1],
        mean_rgb[:, :, :, 2],
    )
    d_rg = tf.square(mr - mg)
    d_rb = tf.square(mr - mb)
    d_gb = tf.square(mb - mg)
    return tf.sqrt(tf.square(d_rg) + tf.square(d_rb) + tf.square(d_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





class SpatialConsistencyLoss(keras.losses.Loss):
    def __init__(self, **kwargs):
        super().__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





class ZeroDCE(keras.Model):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.dce_model = build_dce_net()

    def compile(self, learning_rate, **kwargs):
        super().compile(**kwargs)
        self.optimizer = keras.optimizers.Adam(learning_rate=learning_rate)
        self.spatial_constancy_loss = SpatialConsistencyLoss(reduction="none")
        self.total_loss_tracker = keras.metrics.Mean(name="total_loss")
        self.illumination_smoothness_loss_tracker = keras.metrics.Mean(
            name="illumination_smoothness_loss"
        )
        self.spatial_constancy_loss_tracker = keras.metrics.Mean(
            name="spatial_constancy_loss"
        )
        self.color_constancy_loss_tracker = keras.metrics.Mean(
            name="color_constancy_loss"
        )
        self.exposure_loss_tracker = keras.metrics.Mean(name="exposure_loss")

    @property
    def metrics(self):
        return [
            self.total_loss_tracker,
            self.illumination_smoothness_loss_tracker,
            self.spatial_constancy_loss_tracker,
            self.color_constancy_loss_tracker,
            self.exposure_loss_tracker,
        ]

    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)
        )

        self.total_loss_tracker.update_state(losses["total_loss"])
        self.illumination_smoothness_loss_tracker.update_state(
            losses["illumination_smoothness_loss"]
        )
        self.spatial_constancy_loss_tracker.update_state(
            losses["spatial_constancy_loss"]
        )
        self.color_constancy_loss_tracker.update_state(
            losses["color_constancy_loss"]
        )
        self.exposure_loss_tracker.update_state(losses["exposure_loss"])

        return {metric.name: metric.result() for metric in self.metrics}

    def test_step(self, data):
        output = self.dce_model(data)
        losses = self.compute_losses(data, output)

        self.total_loss_tracker.update_state(losses["total_loss"])
        self.illumination_smoothness_loss_tracker.update_state(
            losses["illumination_smoothness_loss"]
        )
        self.spatial_constancy_loss_tracker.update_state(
            losses["spatial_constancy_loss"]
        )
        self.color_constancy_loss_tracker.update_state(
            losses["color_constancy_loss"]
        )
        self.exposure_loss_tracker.update_state(losses["exposure_loss"])

        return {metric.name: metric.result() for metric in self.metrics}

    def save_weights(
        self, filepath, overwrite=True, save_format=None, options=None
    ):
        
        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
    ):
        
        self.dce_model.load_weights(
            filepath=filepath,
            by_name=by_name,
            skip_mismatch=skip_mismatch,
            options=options,
        )




zero_dce_model = ZeroDCE()
zero_dce_model.compile(learning_rate=1e-4)
history = zero_dce_model.fit(
    train_dataset, validation_data=val_dataset, epochs=100
)


def plot_result(item):
    plt.plot(history.history[item], label=item)
    plt.plot(history.history["val_" + item], label="val_" + item)
    plt.xlabel("Epochs")
    plt.ylabel(item)
    plt.title("Train and Validation {} Over Epochs".format(item), fontsize=14)
    plt.legend()
    plt.grid()
    plt.show()


plot_result("total_loss")
plot_result("illumination_smoothness_loss")
plot_result("spatial_constancy_loss")
plot_result("color_constancy_loss")
plot_result("exposure_loss")




def plot_results(images, titles, figure_size=(12, 12)):
    fig = plt.figure(figsize=figure_size)
    for i in range(len(images)):
        fig.add_subplot(1, len(images), i + 1).set_title(titles[i])
        _ = plt.imshow(images[i])
        plt.axis("off")
    plt.show()


def infer(original_image):
    image = keras.utils.img_to_array(original_image)
    image = image.astype("float32") / 255.0
    image = np.expand_dims(image, axis=0)
    output_image = zero_dce_model(image)
    output_image = tf.cast((output_image[0, :, :, :] * 255), dtype=np.uint8)
    output_image = Image.fromarray(output_image.numpy())
    return output_image




for val_image_file in test_low_light_images:
    original_image = Image.open(val_image_file)
    enhanced_image = infer(original_image)
    plot_results(
        [original_image, ImageOps.autocontrast(original_image), enhanced_image],
        ["Original", "PIL Autocontrast", "Enhanced"],
        (20, 12),
    )