title lowlight image enhancement using mirnet soumik rakshithttp github comsoumik12345 converted to keras 3 by soumik rakshithttp github comsoumik12345 date created 20210911 last modified 20230715 description implementing the mirnet architecture for lowlight image enhancement accelerator gpu introduction with the goal of recovering highquality image content from its degraded version image restoration enjoys numerous applications such as in photography security medical imaging and remote sensing in this example we implement the mirnet model for lowlight image enhancement a fullyconvolutional architecture that learns an enriched set of features that combines contextual information from multiple scales while simultaneously preserving the highresolution spatial details references learning enriched features for real image restoration and enhancementhttps arxiv orgabs2003 06792 the retinex theory of color visionhttp www cnbc cmu edutaicppaperse landretinextheoryscientifcamerican pdf two deterministic halfquadratic regularization algorithms for computed imaginghttps ieeexplore ieee orgdocument413553 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 pip install q githttps github comkerasteamkeras shell wget https huggingface codatasetsgeekyrakshitloldatasetresolvemainloldataset zip unzip q loldataset zip rm loldataset zip creating a tensorflow dataset we use 300 image pairs from the lol dataset s training set for training and we use the remaining 185 image pairs for validation we generate random crops of size 128 x 128 from the image pairs to be used for both training and validation in order to avoid none during shape inference mirnet model here are the main features of the mirnet model a feature extraction model that computes a complementary set of features across multiple spatial scales while maintaining the original highresolution features to preserve precise spatial details a regularly repeated mechanism for information exchange where the features across multiresolution branches are progressively fused together for improved representation learning a new approach to fuse multiscale features using a selective kernel network that dynamically combines variable receptive fields and faithfully preserves the original feature information at each spatial resolution a recursive residual design that progressively breaks down the input signal in order to simplify the overall learning process and allows the construction of very deep networks https raw githubusercontent comsoumik12345mirnetmasterassetsmirnetarchitecture png selective kernel feature fusion the selective kernel feature fusion or skff module performs dynamic adjustment of receptive fields via two operations fuse and select the fuse operator generates global feature descriptors by combining the information from multiresolution streams the select operator uses these descriptors to recalibrate the feature maps of different streams followed by their aggregation fuse the skff receives inputs from three parallel convolution streams carrying different scales of information we first combine these multiscale features using an elementwise sum on which we apply global average pooling gap across the spatial dimension next we apply a channel downscaling convolution layer to generate a compact feature representation which passes through three parallel channelupscaling convolution layers one for each resolution stream and provides us with three feature descriptors select this operator applies the softmax function to the feature descriptors to obtain the corresponding activations that are used to adaptively recalibrate multiscale feature maps the aggregated features are defined as the sum of product of the corresponding multiscale feature and the feature descriptor https i imgur com7u6ixf6 png dual attention unit the dual attention unit or dau is used to extract features in the convolutional streams while the skff block fuses information across multiresolution branches we also need a mechanism to share information within a feature tensor both along the spatial and the channel dimensions which is done by the dau block the dau suppresses less useful features and only allows more informative ones to pass further this feature recalibration is achieved by using channel attention and spatial attention mechanisms the channel attention branch exploits the interchannel relationships of the convolutional feature maps by applying squeeze and excitation operations given a feature map the squeeze operation applies global average pooling across spatial dimensions to encode global context thus yielding a feature descriptor the excitation operator passes this feature descriptor through two convolutional layers followed by the sigmoid gating and generates activations finally the output of channel attention branch is obtained by rescaling the input feature map with the output activations the spatial attention branch is designed to exploit the interspatial dependencies of convolutional features the goal of spatial attention is to generate a spatial attention map and use it to recalibrate the incoming features to generate the spatial attention map the spatial attention branch first independently applies global average pooling and max pooling operations on input features along the channel dimensions and concatenates the outputs to form a resultant feature map which is then passed through a convolution and sigmoid activation to obtain the spatial attention map this spatial attention map is then used to rescale the input feature map https i imgur comdl0iwqs png multiscale residual block the multiscale residual block is capable of generating a spatiallyprecise output by maintaining highresolution representations while receiving rich contextual information from lowresolutions the mrb consists of multiple three in this paper fullyconvolutional streams connected in parallel it allows information exchange across parallel streams in order to consolidate the highresolution features with the help of lowresolution features and vice versa the mirnet employs a recursive residual design with skip connections to ease the flow of information during the learning process in order to maintain the residual nature of our architecture residual resizing modules are used to perform downsampling and upsampling operations that are used in the multiscale residual block https i imgur comwzzkv57 png recursive residual modules mrb block features dau skff dau 2 skff 2 mirnet model training we train mirnet using charbonnier loss as the loss function and adam optimizer with a learning rate of 1e4 we use peak signal noise ratio or psnr as a metric which is an expression for the ratio between the maximum possible value power of a signal and the power of distorting noise that affects the quality of its representation 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 cokerasiolowlightenhancemirnet and try the demo on hugging face spaceshttps huggingface cospaceskerasioenhancelowlightimage introduction with the goal of recovering high quality image content from its degraded version image restoration enjoys numerous applications such as in photography security medical imaging and remote sensing in this example we implement the mirnet model for low light image enhancement a fully convolutional architecture that learns an enriched set of features that combines contextual information from multiple scales while simultaneously preserving the high resolution spatial details references learning enriched features for real image restoration and enhancement https arxiv org abs 2003 06792 the retinex theory of color vision http www cnbc cmu edu tai cp_papers e land_retinex_theory_scientifcamerican pdf two deterministic half quadratic regularization algorithms for computed imaging https ieeexplore ieee org document 413553 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 pip install q git https github com keras team keras 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 image pairs from the lol dataset s training set for training and we use the remaining 185 image pairs for validation we generate random crops of size 128 x 128 from the image pairs to be used for both training and validation in order to avoid none during shape inference mirnet model here are the main features of the mirnet model a feature extraction model that computes a complementary set of features across multiple spatial scales while maintaining the original high resolution features to preserve precise spatial details a regularly repeated mechanism for information exchange where the features across multi resolution branches are progressively fused together for improved representation learning a new approach to fuse multi scale features using a selective kernel network that dynamically combines variable receptive fields and faithfully preserves the original feature information at each spatial resolution a recursive residual design that progressively breaks down the input signal in order to simplify the overall learning process and allows the construction of very deep networks https raw githubusercontent com soumik12345 mirnet master assets mirnet_architecture png selective kernel feature fusion the selective kernel feature fusion or skff module performs dynamic adjustment of receptive fields via two operations fuse and select the fuse operator generates global feature descriptors by combining the information from multi resolution streams the select operator uses these descriptors to recalibrate the feature maps of different streams followed by their aggregation fuse the skff receives inputs from three parallel convolution streams carrying different scales of information we first combine these multi scale features using an element wise sum on which we apply global average pooling gap across the spatial dimension next we apply a channel downscaling convolution layer to generate a compact feature representation which passes through three parallel channel upscaling convolution layers one for each resolution stream and provides us with three feature descriptors select this operator applies the softmax function to the feature descriptors to obtain the corresponding activations that are used to adaptively recalibrate multi scale feature maps the aggregated features are defined as the sum of product of the corresponding multi scale feature and the feature descriptor https i imgur com 7u6ixf6 png dual attention unit the dual attention unit or dau is used to extract features in the convolutional streams while the skff block fuses information across multi resolution branches we also need a mechanism to share information within a feature tensor both along the spatial and the channel dimensions which is done by the dau block the dau suppresses less useful features and only allows more informative ones to pass further this feature recalibration is achieved by using channel attention and spatial attention mechanisms the channel attention branch exploits the inter channel relationships of the convolutional feature maps by applying squeeze and excitation operations given a feature map the squeeze operation applies global average pooling across spatial dimensions to encode global context thus yielding a feature descriptor the excitation operator passes this feature descriptor through two convolutional layers followed by the sigmoid gating and generates activations finally the output of channel attention branch is obtained by rescaling the input feature map with the output activations the spatial attention branch is designed to exploit the inter spatial dependencies of convolutional features the goal of spatial attention is to generate a spatial attention map and use it to recalibrate the incoming features to generate the spatial attention map the spatial attention branch first independently applies global average pooling and max pooling operations on input features along the channel dimensions and concatenates the outputs to form a resultant feature map which is then passed through a convolution and sigmoid activation to obtain the spatial attention map this spatial attention map is then used to rescale the input feature map https i imgur com dl0iwqs png multi scale residual block the multi scale residual block is capable of generating a spatially precise output by maintaining high resolution representations while receiving rich contextual information from low resolutions the mrb consists of multiple three in this paper fully convolutional streams connected in parallel it allows information exchange across parallel streams in order to consolidate the high resolution features with the help of low resolution features and vice versa the mirnet employs a recursive residual design with skip connections to ease the flow of information during the learning process in order to maintain the residual nature of our architecture residual resizing modules are used to perform downsampling and upsampling operations that are used in the multi scale residual block https i imgur com wzzkv57 png recursive residual modules mrb block features dau skff dau 2 skff 2 mirnet model training we train mirnet using charbonnier loss as the loss function and adam optimizer with a learning rate of 1e 4 we use peak signal noise ratio or psnr as a metric which is an expression for the ratio between the maximum possible value power of a signal and the power of distorting noise that affects the quality of its representation 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 lowlight enhance mirnet and try the demo on hugging face spaces https huggingface co spaces keras io enhance_low_light_image
import os

os.environ["KERAS_BACKEND"] = "tensorflow"

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





random.seed(10)

IMAGE_SIZE = 128
BATCH_SIZE = 4
MAX_TRAIN_IMAGES = 300


def read_image(image_path):
    image = tf.io.read_file(image_path)
    image = tf.image.decode_png(image, channels=3)
    image.set_shape([None, None, 3])
    image = tf.cast(image, dtype=tf.float32) / 255.0
    return image


def random_crop(low_image, enhanced_image):
    low_image_shape = tf.shape(low_image)[:2]
    low_w = tf.random.uniform(
        shape=(), maxval=low_image_shape[1] - IMAGE_SIZE + 1, dtype=tf.int32
    )
    low_h = tf.random.uniform(
        shape=(), maxval=low_image_shape[0] - IMAGE_SIZE + 1, dtype=tf.int32
    )
    low_image_cropped = low_image[
        low_h : low_h + IMAGE_SIZE, low_w : low_w + IMAGE_SIZE
    ]
    enhanced_image_cropped = enhanced_image[
        low_h : low_h + IMAGE_SIZE, low_w : low_w + IMAGE_SIZE
    ]
    
    low_image_cropped.set_shape([IMAGE_SIZE, IMAGE_SIZE, 3])
    enhanced_image_cropped.set_shape([IMAGE_SIZE, IMAGE_SIZE, 3])
    return low_image_cropped, enhanced_image_cropped


def load_data(low_light_image_path, enhanced_image_path):
    low_light_image = read_image(low_light_image_path)
    enhanced_image = read_image(enhanced_image_path)
    low_light_image, enhanced_image = random_crop(
        low_light_image, enhanced_image
    )
    return low_light_image, enhanced_image


def get_dataset(low_light_images, enhanced_images):
    dataset = tf.data.Dataset.from_tensor_slices(
        (low_light_images, enhanced_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
]
train_enhanced_images = sorted(glob("./lol_dataset/our485/high/*"))[
    :MAX_TRAIN_IMAGES
]

val_low_light_images = sorted(glob("./lol_dataset/our485/low/*"))[
    MAX_TRAIN_IMAGES:
]
val_enhanced_images = sorted(glob("./lol_dataset/our485/high/*"))[
    MAX_TRAIN_IMAGES:
]

test_low_light_images = sorted(glob("./lol_dataset/eval15/low/*"))
test_enhanced_images = sorted(glob("./lol_dataset/eval15/high/*"))


train_dataset = get_dataset(train_low_light_images, train_enhanced_images)
val_dataset = get_dataset(val_low_light_images, val_enhanced_images)


print("Train Dataset:", train_dataset.element_spec)
print("Val Dataset:", val_dataset.element_spec)






def selective_kernel_feature_fusion(
    multi_scale_feature_1, multi_scale_feature_2, multi_scale_feature_3
):
    channels = list(multi_scale_feature_1.shape)[-1]
    combined_feature = layers.Add()(
        [multi_scale_feature_1, multi_scale_feature_2, multi_scale_feature_3]
    )
    gap = layers.GlobalAveragePooling2D()(combined_feature)
    channel_wise_statistics = layers.Reshape((1, 1, channels))(gap)
    compact_feature_representation = layers.Conv2D(
        filters=channels // 8, kernel_size=(1, 1), activation="relu"
    )(channel_wise_statistics)
    feature_descriptor_1 = layers.Conv2D(
        channels, kernel_size=(1, 1), activation="softmax"
    )(compact_feature_representation)
    feature_descriptor_2 = layers.Conv2D(
        channels, kernel_size=(1, 1), activation="softmax"
    )(compact_feature_representation)
    feature_descriptor_3 = layers.Conv2D(
        channels, kernel_size=(1, 1), activation="softmax"
    )(compact_feature_representation)
    feature_1 = multi_scale_feature_1 * feature_descriptor_1
    feature_2 = multi_scale_feature_2 * feature_descriptor_2
    feature_3 = multi_scale_feature_3 * feature_descriptor_3
    aggregated_feature = layers.Add()([feature_1, feature_2, feature_3])
    return aggregated_feature





class ChannelPooling(layers.Layer):
    def __init__(self, axis=-1, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.axis = axis
        self.concat = layers.Concatenate(axis=self.axis)

    def call(self, inputs):
        average_pooling = tf.expand_dims(
            tf.reduce_mean(inputs, axis=-1), axis=-1
        )
        max_pooling = tf.expand_dims(tf.reduce_max(inputs, axis=-1), axis=-1)
        return self.concat([average_pooling, max_pooling])

    def get_config(self):
        config = super().get_config()
        config.update({"axis": self.axis})


def spatial_attention_block(input_tensor):
    compressed_feature_map = ChannelPooling(axis=-1)(input_tensor)
    feature_map = layers.Conv2D(1, kernel_size=(1, 1))(compressed_feature_map)
    feature_map = keras.activations.sigmoid(feature_map)
    return input_tensor * feature_map


def channel_attention_block(input_tensor):
    channels = list(input_tensor.shape)[-1]
    average_pooling = layers.GlobalAveragePooling2D()(input_tensor)
    feature_descriptor = layers.Reshape((1, 1, channels))(average_pooling)
    feature_activations = layers.Conv2D(
        filters=channels // 8, kernel_size=(1, 1), activation="relu"
    )(feature_descriptor)
    feature_activations = layers.Conv2D(
        filters=channels, kernel_size=(1, 1), activation="sigmoid"
    )(feature_activations)
    return input_tensor * feature_activations


def dual_attention_unit_block(input_tensor):
    channels = list(input_tensor.shape)[-1]
    feature_map = layers.Conv2D(
        channels, kernel_size=(3, 3), padding="same", activation="relu"
    )(input_tensor)
    feature_map = layers.Conv2D(channels, kernel_size=(3, 3), padding="same")(
        feature_map
    )
    channel_attention = channel_attention_block(feature_map)
    spatial_attention = spatial_attention_block(feature_map)
    concatenation = layers.Concatenate(axis=-1)(
        [channel_attention, spatial_attention]
    )
    concatenation = layers.Conv2D(channels, kernel_size=(1, 1))(concatenation)
    return layers.Add()([input_tensor, concatenation])







def down_sampling_module(input_tensor):
    channels = list(input_tensor.shape)[-1]
    main_branch = layers.Conv2D(
        channels, kernel_size=(1, 1), activation="relu"
    )(input_tensor)
    main_branch = layers.Conv2D(
        channels, kernel_size=(3, 3), padding="same", activation="relu"
    )(main_branch)
    main_branch = layers.MaxPooling2D()(main_branch)
    main_branch = layers.Conv2D(channels * 2, kernel_size=(1, 1))(main_branch)
    skip_branch = layers.MaxPooling2D()(input_tensor)
    skip_branch = layers.Conv2D(channels * 2, kernel_size=(1, 1))(skip_branch)
    return layers.Add()([skip_branch, main_branch])


def up_sampling_module(input_tensor):
    channels = list(input_tensor.shape)[-1]
    main_branch = layers.Conv2D(
        channels, kernel_size=(1, 1), activation="relu"
    )(input_tensor)
    main_branch = layers.Conv2D(
        channels, kernel_size=(3, 3), padding="same", activation="relu"
    )(main_branch)
    main_branch = layers.UpSampling2D()(main_branch)
    main_branch = layers.Conv2D(channels // 2, kernel_size=(1, 1))(main_branch)
    skip_branch = layers.UpSampling2D()(input_tensor)
    skip_branch = layers.Conv2D(channels // 2, kernel_size=(1, 1))(skip_branch)
    return layers.Add()([skip_branch, main_branch])



def multi_scale_residual_block(input_tensor, channels):
    
    level1 = input_tensor
    level2 = down_sampling_module(input_tensor)
    level3 = down_sampling_module(level2)
    
    level1_dau = dual_attention_unit_block(level1)
    level2_dau = dual_attention_unit_block(level2)
    level3_dau = dual_attention_unit_block(level3)
    
    level1_skff = selective_kernel_feature_fusion(
        level1_dau,
        up_sampling_module(level2_dau),
        up_sampling_module(up_sampling_module(level3_dau)),
    )
    level2_skff = selective_kernel_feature_fusion(
        down_sampling_module(level1_dau),
        level2_dau,
        up_sampling_module(level3_dau),
    )
    level3_skff = selective_kernel_feature_fusion(
        down_sampling_module(down_sampling_module(level1_dau)),
        down_sampling_module(level2_dau),
        level3_dau,
    )
    
    level1_dau_2 = dual_attention_unit_block(level1_skff)
    level2_dau_2 = up_sampling_module((dual_attention_unit_block(level2_skff)))
    level3_dau_2 = up_sampling_module(
        up_sampling_module(dual_attention_unit_block(level3_skff))
    )
    
    skff_ = selective_kernel_feature_fusion(
        level1_dau_2, level2_dau_2, level3_dau_2
    )
    conv = layers.Conv2D(channels, kernel_size=(3, 3), padding="same")(skff_)
    return layers.Add()([input_tensor, conv])





def recursive_residual_group(input_tensor, num_mrb, channels):
    conv1 = layers.Conv2D(channels, kernel_size=(3, 3), padding="same")(
        input_tensor
    )
    for _ in range(num_mrb):
        conv1 = multi_scale_residual_block(conv1, channels)
    conv2 = layers.Conv2D(channels, kernel_size=(3, 3), padding="same")(conv1)
    return layers.Add()([conv2, input_tensor])


def mirnet_model(num_rrg, num_mrb, channels):
    input_tensor = keras.Input(shape=[None, None, 3])
    x1 = layers.Conv2D(channels, kernel_size=(3, 3), padding="same")(
        input_tensor
    )
    for _ in range(num_rrg):
        x1 = recursive_residual_group(x1, num_mrb, channels)
    conv = layers.Conv2D(3, kernel_size=(3, 3), padding="same")(x1)
    output_tensor = layers.Add()([input_tensor, conv])
    return keras.Model(input_tensor, output_tensor)


model = mirnet_model(num_rrg=3, num_mrb=2, channels=64)




def charbonnier_loss(y_true, y_pred):
    return tf.reduce_mean(tf.sqrt(tf.square(y_true - y_pred) + tf.square(1e-3)))


def peak_signal_noise_ratio(y_true, y_pred):
    return tf.image.psnr(y_pred, y_true, max_val=255.0)


optimizer = keras.optimizers.Adam(learning_rate=1e-4)
model.compile(
    optimizer=optimizer,
    loss=charbonnier_loss,
    metrics=[peak_signal_noise_ratio],
)

history = model.fit(
    train_dataset,
    validation_data=val_dataset,
    epochs=50,
    callbacks=[
        keras.callbacks.ReduceLROnPlateau(
            monitor="val_peak_signal_noise_ratio",
            factor=0.5,
            patience=5,
            verbose=1,
            min_delta=1e-7,
            mode="max",
        )
    ],
)


def plot_history(value):
    plt.plot(history.history[value], label=f"train_{value}")
    plt.plot(history.history[f"val_{value}"], label=f"val_{value}")
    plt.xlabel("Epochs")
    plt.ylabel(value)
    plt.title(f"Train and Validation {value} Over Epochs", fontsize=14)
    plt.legend()
    plt.grid()
    plt.show()


plot_history("loss")
plot_history("peak_signal_noise_ratio")




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 = model.predict(image, verbose=0)
    output_image = output[0] * 255.0
    output_image = output_image.clip(0, 255)
    output_image = output_image.reshape(
        (np.shape(output_image)[0], np.shape(output_image)[1], 3)
    )
    output_image = Image.fromarray(np.uint8(output_image))
    original_image = Image.fromarray(np.uint8(original_image))
    return output_image





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