added resize option

enhance_me/zero_dce/dataloader.py

@@ -9,20 +9,31 @@ class UnpairedLowLightDataset:
     def __init__(
         self,
         image_size: int = 256,
+        apply_resize: bool = False,
         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.image_size = image_size
+        self.apply_resize = apply_resize
         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 _resize(self, image):
+        return tf.image.resize(image, (self.image_size, self.image_size))
+
     def _get_dataset(self, images: List[str], batch_size: int, is_train: bool):
         dataset = tf.data.Dataset.from_tensor_slices((images))
         dataset = dataset.map(read_image, num_parallel_calls=tf.data.AUTOTUNE)
-        dataset = dataset.map(
-            self.augmentation_factory.random_crop, num_parallel_calls=tf.data.AUTOTUNE
+        dataset = (
+            dataset.map(
+                self.augmentation_factory.random_crop,
+                num_parallel_calls=tf.data.AUTOTUNE,
+            )
+            if not self.apply_resize
+            else dataset.map(self._resize, num_parallel_calls=tf.data.AUTOTUNE)
         )
         if is_train:
             dataset = (

enhance_me/zero_dce/zero_dce.py

@@ -113,6 +113,7 @@ class ZeroDCE(Model):
         self,
         image_size: int = 256,
         dataset_label: str = "lol",
+        apply_resize: bool = False,
         apply_random_horizontal_flip: bool = True,
         apply_random_vertical_flip: bool = True,
         apply_random_rotation: bool = True,
@@ -123,6 +124,7 @@ class ZeroDCE(Model):
             (self.low_images, _), (self.test_low_images, _) = download_lol_dataset()
         data_loader = UnpairedLowLightDataset(
             image_size,
+            apply_resize,
             apply_random_horizontal_flip,
             apply_random_vertical_flip,
             apply_random_rotation,
@@ -130,7 +132,7 @@ class ZeroDCE(Model):
         self.train_dataset, self.val_dataset = data_loader.get_datasets(
             self.low_images, val_split, batch_size
         )
-    
+
     def train(self, epochs: int):
         log_dir = os.path.join(
             self.experiment_name,
@@ -148,7 +150,7 @@ class ZeroDCE(Model):
             callbacks=callbacks,
         )
         return history
-    
+
     def infer(self, original_image):
         image = keras.preprocessing.image.img_to_array(original_image)
         image = image.astype("float32") / 255.0
@@ -157,7 +159,7 @@ class ZeroDCE(Model):
         output_image = tf.cast((output_image[0, :, :, :] * 255), dtype=np.uint8)
         output_image = Image.fromarray(output_image.numpy())
         return output_image
-    
+
     def infer_from_file(self, original_image_file: str):
         original_image = Image.open(original_image_file)
         return self.infer(original_image)

notebooks/enhance_me_train.ipynb

@@ -190,6 +190,7 @@
     "experiment_name = \"lol_dataset_256\"  # @param {type:\"string\"}\n",
     "image_size = 256  # @param {type:\"integer\"}\n",
     "dataset_label = \"lol\"  # @param [\"lol\"]\n",
+    "apply_resize = False  # @param {type:\"boolean\"}\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",
@@ -223,6 +224,7 @@
     "zero_dce.build_datasets(\n",
     "    image_size=image_size,\n",
     "    dataset_label=dataset_label,\n",
+    "    apply_resize=apply_resize,\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",