added download function for lol dataset

.gitignore

@@ -127,3 +127,6 @@ dmypy.json
 
 # Pyre type checker
 .pyre/
+
+# Datasets
+datasets/

enhance_me/commons.py

@@ -1,8 +1,11 @@
 import os
 import wandb
-import tensorflow as tf
+from glob import glob
 import matplotlib.pyplot as plt
 
+import tensorflow as tf
+from tensorflow.keras import utils
+
 
 def read_image(image_path):
     image = tf.io.read_file(image_path)
@@ -39,5 +42,22 @@ def closest_number(n, m):
 
 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
+        os.environ["WANDB_API_KEY"] = wandb_api_key
         wandb.init(project=project_name, name=experiment_name)
+
+
+def download_lol_dataset():
+    utils.get_file(
+        "lol_dataset.zip",
+        "https://github.com/soumik12345/enhance-me/releases/download/v0.1/lol_dataset.zip",
+        cache_dir="./",
+        cache_subdir="./datasets",
+        extract=True,
+    )
+    low_images = sorted(glob("./datasets/lol_dataset/our485/low/*"))
+    enhanced_images = sorted(glob("./datasets/lol_dataset/our485/high/*"))
+    assert len(low_images) == len(enhanced_images)
+    test_low_images = sorted(glob("./datasets/lol_dataset/eval15/low/*"))
+    test_enhanced_images = sorted(glob("./datasets/lol_dataset/eval15/high/*"))
+    assert len(test_low_images) == len(test_enhanced_images)
+    return (low_images, enhanced_images), (test_low_images, test_enhanced_images)

enhance_me/mirnet/mirnet.py

@@ -12,7 +12,12 @@ 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
+from ..commons import (
+    peak_signal_noise_ratio,
+    closest_number,
+    init_wandb,
+    download_lol_dataset,
+)
 
 
 class MIRNet:
@@ -20,12 +25,15 @@ class MIRNet:
         self,
         experiment_name: str,
         image_size: int = 256,
+        dataset_label: str = "lol",
         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
+        if dataset_label == "lol":
+            download_lol_dataset()
         self.data_loader = LowLightDataset(
             image_size=image_size,
             apply_random_horizontal_flip=apply_random_horizontal_flip,