Merge pull request #1 from soumik12345/mirnet

.gitignore

@@ -127,3 +127,8 @@ dmypy.json
 
 # Pyre type checker
 .pyre/
+
+# Datasets
+datasets/
+**.zip
+**.h5

README.md

@@ -1 +1,11 @@
+---
+title: Enhance Me
+emoji: 🌖
+colorFrom: pink
+colorTo: pink
+sdk: streamlit
+app_file: app.py
+pinned: false
+---
+
 # enhance-me

app.py

@@ -0,0 +1,40 @@
+from PIL import Image
+import streamlit as st
+from tensorflow.keras import utils
+
+from enhance_me.mirnet import MIRNet
+
+
+@st.cache
+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.load_weights("./weights/weights_lol_128.h5")
+    return mirnet
+
+
+def main():
+    st.markdown("# Enhance Me")
+    st.markdown("Made with :heart: by [geekyRakshit](http://github.com/soumik12345)")
+    application = st.sidebar.selectbox(
+        "Please select the application:", ("", "Low-light enhancement")
+    )
+    if application != "":
+        if application == "Low-light enhancement":
+            uploaded_file = st.sidebar.file_uploader("Select your image:")
+            if uploaded_file is not None:
+                original_image = Image.open(uploaded_file)
+                st.image(original_image, caption="original image")
+                mirnet = get_mirnet_object()
+                enhanced_image = mirnet.infer(original_image)
+                st.image(enhanced_image, caption="enhanced image")
+
+
+if __name__ == "__main__":
+    main()

enhance_me/augmentation.py

@@ -0,0 +1,51 @@
+import tensorflow as tf
+
+
+class AugmentationFactory:
+    def __init__(self, image_size) -> None:
+        self.image_size = image_size
+
+    def random_crop(self, input_image, enhanced_image):
+        input_image_shape = tf.shape(input_image)[:2]
+        low_w = tf.random.uniform(
+            shape=(), maxval=input_image_shape[1] - self.image_size + 1, dtype=tf.int32
+        )
+        low_h = tf.random.uniform(
+            shape=(), maxval=input_image_shape[0] - self.image_size + 1, dtype=tf.int32
+        )
+        enhanced_w = low_w
+        enhanced_h = low_h
+        input_image_cropped = input_image[
+            low_h : low_h + self.image_size, low_w : low_w + self.image_size
+        ]
+        enhanced_image_cropped = enhanced_image[
+            enhanced_h : enhanced_h + self.image_size,
+            enhanced_w : enhanced_w + self.image_size,
+        ]
+        return input_image_cropped, enhanced_image_cropped
+
+    def random_horizontal_flip(sefl, input_image, enhanced_image):
+        return tf.cond(
+            tf.random.uniform(shape=(), maxval=1) < 0.5,
+            lambda: (input_image, enhanced_image),
+            lambda: (
+                tf.image.flip_left_right(input_image),
+                tf.image.flip_left_right(enhanced_image),
+            ),
+        )
+
+    def random_vertical_flip(self, input_image, enhanced_image):
+        return tf.cond(
+            tf.random.uniform(shape=(), maxval=1) < 0.5,
+            lambda: (input_image, enhanced_image),
+            lambda: (
+                tf.image.flip_up_down(input_image),
+                tf.image.flip_up_down(enhanced_image),
+            ),
+        )
+
+    def random_rotate(self, input_image, enhanced_image):
+        condition = tf.random.uniform(shape=(), maxval=4, dtype=tf.int32)
+        return tf.image.rot90(input_image, condition), tf.image.rot90(
+            enhanced_image, condition
+        )

enhance_me/commons.py

@@ -0,0 +1,63 @@
+import os
+import wandb
+from glob import glob
+import matplotlib.pyplot as plt
+
+import tensorflow as tf
+from tensorflow.keras import utils
+
+
+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 peak_signal_noise_ratio(y_true, y_pred):
+    return tf.image.psnr(y_pred, y_true, max_val=255.0)
+
+
+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 closest_number(n, m):
+    q = int(n / m)
+    n1 = m * q
+    if (n * m) > 0:
+        n2 = m * (q + 1)
+    else:
+        n2 = m * (q - 1)
+    if abs(n - n1) < abs(n - n2):
+        return n1
+    return n2
+
+
+def init_wandb(project_name, experiment_name, wandb_api_key):
+    if project_name is not None and experiment_name is not None:
+        os.environ["WANDB_API_KEY"] = wandb_api_key
+        wandb.init(project=project_name, name=experiment_name, sync_tensorboard=True)
+
+
+def download_lol_dataset():
+    utils.get_file(
+        "lol_dataset.zip",
+        "https://github.com/soumik12345/enhance-me/releases/download/v0.1/lol_dataset.zip",
+        cache_dir="./",
+        cache_subdir="./datasets",
+        extract=True,
+    )
+    low_images = sorted(glob("./datasets/lol_dataset/our485/low/*"))
+    enhanced_images = sorted(glob("./datasets/lol_dataset/our485/high/*"))
+    assert len(low_images) == len(enhanced_images)
+    test_low_images = sorted(glob("./datasets/lol_dataset/eval15/low/*"))
+    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)

enhance_me/mirnet/__init__.py

@@ -0,0 +1 @@
+from .mirnet import MIRNet

enhance_me/mirnet/dataloader.py

@@ -0,0 +1,92 @@
+import tensorflow as tf
+from typing import List
+
+from ..commons import read_image
+from ..augmentation import AugmentationFactory
+
+
+class LowLightDataset:
+    def __init__(
+        self,
+        image_size: int = 256,
+        apply_random_horizontal_flip: bool = True,
+        apply_random_vertical_flip: bool = True,
+        apply_random_rotation: bool = True,
+    ) -> None:
+        self.augmentation_factory = AugmentationFactory(image_size=image_size)
+        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 load_data(self, 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 = self.augmentation_factory.random_crop(
+            low_light_image, enhanced_image
+        )
+        return low_light_image, enhanced_image
+
+    def _get_dataset(
+        self,
+        low_light_images: List[str],
+        enhanced_images: List[str],
+        batch_size: int = 16,
+        is_train: bool = True,
+    ):
+        dataset = tf.data.Dataset.from_tensor_slices(
+            (low_light_images, enhanced_images)
+        )
+        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
+        )
+        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,
+        low_light_images: List[str],
+        enhanced_images: List[str],
+        val_split: float = 0.2,
+        batch_size: int = 16,
+    ):
+        assert len(low_light_images) == len(enhanced_images)
+        split_index = int(len(low_light_images) * (1 - val_split))
+        train_low_light_images = low_light_images[:split_index]
+        train_enhanced_images = enhanced_images[:split_index]
+        val_low_light_images = low_light_images[split_index:]
+        val_enhanced_images = enhanced_images[split_index:]
+        print(f"Number of train data points: {len(train_low_light_images)}")
+        print(f"Number of validation data points: {len(val_low_light_images)}")
+        train_dataset = self._get_dataset(
+            train_low_light_images, train_enhanced_images, batch_size, is_train=True
+        )
+        val_dataset = self._get_dataset(
+            val_low_light_images, val_enhanced_images, batch_size, is_train=False
+        )
+        return train_dataset, val_dataset

enhance_me/mirnet/losses.py

@@ -0,0 +1,13 @@
+import tensorflow as tf
+from tensorflow.keras import losses
+
+
+class CharbonnierLoss(losses.Loss):
+    def __init__(self, epsilon: float = 1e-3, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.epsilon = epsilon
+
+    def call(self, y_true, y_pred):
+        return tf.reduce_mean(
+            tf.sqrt(tf.square(y_true - y_pred) + tf.square(self.epsilon))
+        )

enhance_me/mirnet/mirnet.py

@@ -0,0 +1,170 @@
+import os
+import numpy as np
+from PIL import Image
+from typing import List
+from datetime import datetime
+
+from tensorflow import keras
+from tensorflow.keras import optimizers, models
+
+from wandb.keras import WandbCallback
+
+from .dataloader import LowLightDataset
+from .models import build_mirnet_model
+from .losses import CharbonnierLoss
+from ..commons import (
+    peak_signal_noise_ratio,
+    closest_number,
+    init_wandb,
+    download_lol_dataset,
+)
+
+
+class MIRNet:
+    def __init__(self, experiment_name=None, wandb_api_key=None) -> None:
+        self.experiment_name = experiment_name
+        if wandb_api_key is not None:
+            init_wandb("mirnet", experiment_name, wandb_api_key)
+            self.using_wandb = True
+        else:
+            self.using_wandb = False
+
+    def build_datasets(
+        self,
+        image_size: int = 256,
+        dataset_label: str = "lol",
+        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,
+    ):
+        if dataset_label == "lol":
+            (self.low_images, self.enhanced_images), (
+                self.test_low_images,
+                self.test_enhanced_images,
+            ) = download_lol_dataset()
+        self.data_loader = LowLightDataset(
+            image_size=image_size,
+            apply_random_horizontal_flip=apply_random_horizontal_flip,
+            apply_random_vertical_flip=apply_random_vertical_flip,
+            apply_random_rotation=apply_random_rotation,
+        )
+        (self.train_dataset, self.val_dataset) = self.data_loader.get_datasets(
+            low_light_images=self.low_images,
+            enhanced_images=self.enhanced_images,
+            val_split=val_split,
+            batch_size=batch_size,
+        )
+
+    def build_model(
+        self,
+        num_recursive_residual_groups: int = 3,
+        num_multi_scale_residual_blocks: int = 2,
+        channels: int = 64,
+        learning_rate: float = 1e-4,
+        epsilon: float = 1e-3,
+    ):
+        self.model = build_mirnet_model(
+            num_rrg=num_recursive_residual_groups,
+            num_mrb=num_multi_scale_residual_blocks,
+            channels=channels,
+        )
+        self.model.compile(
+            optimizer=optimizers.Adam(learning_rate=learning_rate),
+            loss=CharbonnierLoss(epsilon=epsilon),
+            metrics=[peak_signal_noise_ratio],
+        )
+
+    def load_model(
+        self, filepath, custom_objects=None, compile=True, options=None
+    ) -> None:
+        self.model = models.load_model(
+            filepath=filepath,
+            custom_objects=custom_objects,
+            compile=compile,
+            options=options,
+        )
+
+    def save_weights(self, filepath, overwrite=True, save_format=None, options=None):
+        self.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.model.load_weights(
+            filepath, by_name=by_name, skip_mismatch=skip_mismatch, options=options
+        )
+
+    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)
+        model_checkpoint_callback = keras.callbacks.ModelCheckpoint(
+            os.path.join(self.experiment_name, "weights.h5"),
+            save_best_only=True,
+            save_weights_only=True,
+        )
+        reduce_lr_callback = keras.callbacks.ReduceLROnPlateau(
+            monitor="val_peak_signal_noise_ratio",
+            factor=0.5,
+            patience=5,
+            verbose=1,
+            min_delta=1e-7,
+            mode="max",
+        )
+        callbacks = [
+            tensorboard_callback,
+            model_checkpoint_callback,
+            reduce_lr_callback,
+        ]
+        if self.using_wandb:
+            callbacks += [WandbCallback()]
+        history = self.model.fit(
+            self.train_dataset,
+            validation_data=self.val_dataset,
+            epochs=epochs,
+            callbacks=callbacks,
+        )
+        return history
+
+    def infer(
+        self,
+        original_image,
+        image_resize_factor: float = 1.0,
+        resize_output: bool = False,
+    ):
+        width, height = original_image.size
+        target_width, target_height = (
+            closest_number(width // image_resize_factor, 4),
+            closest_number(height // image_resize_factor, 4),
+        )
+        original_image = original_image.resize(
+            (target_width, target_height), Image.ANTIALIAS
+        )
+        image = keras.preprocessing.image.img_to_array(original_image)
+        image = image.astype("float32") / 255.0
+        image = np.expand_dims(image, axis=0)
+        output = self.model.predict(image)
+        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))
+        if resize_output:
+            output_image = output_image.resize((width, height), Image.ANTIALIAS)
+        return output_image
+
+    def infer_from_file(
+        self,
+        original_image_file: str,
+        image_resize_factor: float = 1.0,
+        resize_output: bool = False,
+    ):
+        original_image = Image.open(original_image_file)
+        return self.infer(original_image, image_resize_factor, resize_output)

enhance_me/mirnet/models/__init__.py

@@ -0,0 +1 @@
+from .mirnet_model import build_mirnet_model

enhance_me/mirnet/models/dual_attention.py

@@ -0,0 +1,41 @@
+import tensorflow as tf
+from tensorflow.keras import layers
+
+
+def spatial_attention_block(input_tensor):
+    average_pooling = tf.reduce_max(input_tensor, axis=-1)
+    average_pooling = tf.expand_dims(average_pooling, axis=-1)
+    max_pooling = tf.reduce_mean(input_tensor, axis=-1)
+    max_pooling = tf.expand_dims(max_pooling, axis=-1)
+    concatenated = layers.Concatenate(axis=-1)([average_pooling, max_pooling])
+    feature_map = layers.Conv2D(1, kernel_size=(1, 1))(concatenated)
+    feature_map = tf.nn.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 = tf.reshape(average_pooling, shape=(-1, 1, 1, channels))
+    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])

enhance_me/mirnet/models/mirnet_model.py

@@ -0,0 +1,13 @@
+from tensorflow.keras import layers, Input, Model
+
+from .recursive_residual_blocks import recursive_residual_group
+
+
+def build_mirnet_model(num_rrg, num_mrb, channels):
+    input_tensor = 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 Model(input_tensor, output_tensor)

enhance_me/mirnet/models/recursive_residual_blocks.py

@@ -0,0 +1,79 @@
+import tensorflow as tf
+from tensorflow.keras import layers
+
+from .skff import selective_kernel_feature_fusion
+from .dual_attention import dual_attention_unit_block
+
+
+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])
+
+
+# MRB Block
+def multi_scale_residual_block(input_tensor, channels):
+    # features
+    level1 = input_tensor
+    level2 = down_sampling_module(input_tensor)
+    level3 = down_sampling_module(level2)
+    # DAU
+    level1_dau = dual_attention_unit_block(level1)
+    level2_dau = dual_attention_unit_block(level2)
+    level3_dau = dual_attention_unit_block(level3)
+    # SKFF
+    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,
+    )
+    # DAU 2
+    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 2
+    skff_ = selective_kernel_feature_fusion(level1_dau_2, level3_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])

enhance_me/mirnet/models/skff.py

@@ -0,0 +1,30 @@
+import tensorflow as tf
+from tensorflow.keras import layers
+
+
+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 = tf.reshape(gap, shape=(-1, 1, 1, channels))
+    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

notebooks/enhance_me_train.ipynb

@@ -0,0 +1,207 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "colab_type": "text",
+    "id": "view-in-github"
+   },
+   "source": [
+    "<a href=\"https://colab.research.google.com/github/soumik12345/enhance-me/blob/mirnet/notebooks/enhance_me_train.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/"
+    },
+    "id": "1JryaVhtBHij",
+    "outputId": "97ee6a4a-2479-4124-e96a-f0a792bdec46"
+   },
+   "outputs": [],
+   "source": [
+    "!git clone https://github.com/soumik12345/enhance-me -b mirnet\n",
+    "!pip install -qqq wandb streamlit"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "G_c4VtXWHR5l"
+   },
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import sys\n",
+    "\n",
+    "sys.path.append(\"./enhance-me\")\n",
+    "\n",
+    "from PIL import Image\n",
+    "from enhance_me import commons\n",
+    "from enhance_me.mirnet import MIRNet"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "ZpBHbYaMIqP_"
+   },
+   "outputs": [],
+   "source": [
+    "# @title MIRNet Train Configs\n",
+    "\n",
+    "experiment_name = \"lol_dataset_256\"  # @param {type:\"string\"}\n",
+    "image_size = 128  # @param {type:\"integer\"}\n",
+    "dataset_label = \"lol\"  # @param [\"lol\"]\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 = 4  # @param {type:\"integer\"}\n",
+    "num_recursive_residual_groups = 3  # @param {type:\"slider\", min:1, max:5, step:1}\n",
+    "num_multi_scale_residual_blocks = 2  # @param {type:\"slider\", min:1, max:5, step:1}\n",
+    "learning_rate = 1e-4  # @param {type:\"number\"}\n",
+    "epsilon = 1e-3  # @param {type:\"number\"}\n",
+    "epochs = 50  # @param {type:\"slider\", min:10, max:100, step:5}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/",
+     "height": 52
+    },
+    "id": "IVRoedqBIMuH",
+    "outputId": "53ca5beb-871a-4ec3-b757-173e09a15331"
+   },
+   "outputs": [],
+   "source": [
+    "mirnet = MIRNet(\n",
+    "    experiment_name=experiment_name,\n",
+    "    wandb_api_key=None if wandb_api_key == \"\" else wandb_api_key,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/"
+    },
+    "id": "O66Iwzx8IsGh",
+    "outputId": "0b6f1683-65d1-4737-a32f-d36b331d2bc2"
+   },
+   "outputs": [],
+   "source": [
+    "mirnet.build_datasets(\n",
+    "    image_size=image_size,\n",
+    "    dataset_label=dataset_label,\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": {
+    "id": "tsfKrBCsL_Bb"
+   },
+   "outputs": [],
+   "source": [
+    "mirnet.build_model(\n",
+    "    num_recursive_residual_groups=num_recursive_residual_groups,\n",
+    "    num_multi_scale_residual_blocks=num_multi_scale_residual_blocks,\n",
+    "    learning_rate=learning_rate,\n",
+    "    epsilon=epsilon,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/"
+    },
+    "id": "y3L9wlpkNziL",
+    "outputId": "5149f0e7-91f4-450f-c43a-1b6028692bbc"
+   },
+   "outputs": [],
+   "source": [
+    "history = mirnet.train(epochs=epochs)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "mirnet.load_weights(os.path.join(mirnet.experiment_name, \"weights.h5\"))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "colab": {
+     "background_save": true
+    },
+    "id": "daFKbgBkiyzc"
+   },
+   "outputs": [],
+   "source": [
+    "for index, low_image_file in enumerate(mirnet.test_low_images):\n",
+    "    original_image = Image.open(low_image_file)\n",
+    "    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",
+    "        (18, 18),\n",
+    "    )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "dO-IbNQHkB3R"
+   },
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "accelerator": "GPU",
+  "colab": {
+   "authorship_tag": "ABX9TyN4LuJh6kWhbqxzA5s9sp7k",
+   "collapsed_sections": [],
+   "include_colab_link": true,
+   "machine_shape": "hm",
+   "name": "enhance-me-train.ipynb",
+   "provenance": []
+  },
+  "kernelspec": {
+   "display_name": "Python 3",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
+}

requirements.txt

@@ -0,0 +1,7 @@
+black
+gdown
+matplotlib
+streamlit
+tensorflow
+tqdm
+wandb