added unpaired dataloader

enhance_me/augmentation.py

@@ -49,3 +49,38 @@ class AugmentationFactory:
         return tf.image.rot90(input_image, condition), tf.image.rot90(
             enhanced_image, condition
         )
+
+
+class UnpairedAugmentationFactory:
+    def __init__(self, image_size) -> None:
+        self.image_size = image_size
+
+    def random_crop(self, image):
+        image_shape = tf.shape(image)[:2]
+        crop_w = tf.random.uniform(
+            shape=(), maxval=image_shape[1] - self.image_size + 1, dtype=tf.int32
+        )
+        crop_h = tf.random.uniform(
+            shape=(), maxval=image_shape[0] - self.image_size + 1, dtype=tf.int32
+        )
+        return image[
+            crop_h : crop_h + self.image_size, crop_w : crop_w + self.image_size
+        ]
+
+    def random_horizontal_flip(self, image):
+        return tf.cond(
+            tf.random.uniform(shape=(), maxval=1) < 0.5,
+            lambda: image,
+            lambda: tf.image.flip_left_right(image),
+        )
+
+    def random_vertical_flip(self, image):
+        return tf.cond(
+            tf.random.uniform(shape=(), maxval=1) < 0.5,
+            lambda: image,
+            lambda: tf.image.flip_up_down(image),
+        )
+
+    def random_rotate(self, image):
+        condition = tf.random.uniform(shape=(), maxval=4, dtype=tf.int32)
+        return tf.image.rot90(image, condition)

enhance_me/zero_dce/dataloader.py

@@ -0,0 +1,73 @@
+import tensorflow as tf
+from typing import List
+
+from ..augmentation import UnpairedAugmentationFactory
+
+
+class UnpairedLowLightDataset:
+    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 = UnpairedAugmentationFactory(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, image_path):
+        image = tf.io.read_file(image_path)
+        image = tf.image.decode_png(image, channels=3)
+        image = image / 255.0
+        return image
+
+    def _get_dataset(self, images: List[str], batch_size: int, is_train: bool):
+        dataset = tf.data.Dataset.from_tensor_slices((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,
+        images: List[str],
+        val_split: float = 0.2,
+        batch_size: int = 16,
+    ):
+        split_index = int(len(images) * (1 - val_split))
+        train_images = images[:split_index]
+        val_images = images[split_index:]
+        print(f"Number of train data points: {len(train_images)}")
+        print(f"Number of validation data points: {len(val_images)}")
+        train_dataset = self._get_dataset(train_images, batch_size, is_train=True)
+        val_dataset = self._get_dataset(val_images, batch_size, is_train=False)
+        return train_dataset, val_dataset