updated mirnet

enhance_me/mirnet/mirnet.py

@@ -26,6 +26,7 @@ class MIRNet:
         experiment_name: str,
         image_size: int = 256,
         dataset_label: str = "lol",
+        build_datasets: bool = True,
         val_split: float = 0.2,
         batch_size: int = 16,
         apply_random_horizontal_flip: bool = True,
@@ -39,15 +40,16 @@ class MIRNet:
                 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._build_datasets(
-            low_images, enhanced_images, val_split=val_split, batch_size=batch_size
-        )
+        if build_datasets:
+            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._build_datasets(
+                low_images, enhanced_images, val_split=val_split, batch_size=batch_size
+            )
         if wandb_api_key is not None:
             init_wandb("mirnet", experiment_name, wandb_api_key)
             self.using_wandb = True