added dataloader

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(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,9 @@
+import tensorflow as tf
+
+
+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

enhance_me/mirnet/dataloader.py

@@ -0,0 +1,31 @@
+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) -> None:
+        self.augmentation_factory = AugmentationFactory(image_size=image_size)
+
+    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,
+    ):
+        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.batch(batch_size, drop_remainder=True)
+        return dataset