added mirnet class for training and inference

enhance_me/commons.py

@@ -1,4 +1,7 @@
+import os
+import wandb
 import tensorflow as tf
+import matplotlib.pyplot as plt
 
 
 def read_image(image_path):
@@ -11,3 +14,30 @@ def read_image(image_path):
 
 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)

enhance_me/mirnet/__init__.py

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

enhance_me/mirnet/mirnet.py

@@ -0,0 +1,155 @@
+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
+
+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
+
+
+class MIRNet:
+    def __init__(
+        self,
+        experiment_name: str,
+        image_size: int = 256,
+        apply_random_horizontal_flip: bool = True,
+        apply_random_vertical_flip: bool = True,
+        apply_random_rotation: bool = True,
+        wandb_api_key=None,
+    ) -> None:
+        self.experiment_name = experiment_name
+        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,
+        )
+        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,
+        low_light_images: List[str],
+        enhanced_images: List[str],
+        val_split: float = 0.2,
+        batch_size: int = 16,
+    ):
+        (self.train_dataset, self.val_dataset) = self.data_loader.get_datasets(
+            low_light_images=low_light_images,
+            enhanced_images=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 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.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/mirnet_model.py

@@ -1,10 +1,9 @@
-import tensorflow as tf
 from tensorflow.keras import layers, Input, Model
 
 from .recursive_residual_blocks import recursive_residual_group
 
 
-def mirnet_model(num_rrg, num_mrb, channels):
+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):