updated zero-dce model

enhance_me/zero_dce/__init__.py

@@ -0,0 +1 @@
+from .zero_dce import ZeroDCE

enhance_me/zero_dce/zero_dce.py

@@ -24,7 +24,7 @@ class ZeroDCE(Model):
         super(ZeroDCE, self).__init__(**kwargs)
         self.experiment_name = experiment_name
         if wandb_api_key is not None:
-            init_wandb("mirnet", experiment_name, wandb_api_key)
+            init_wandb("zero-dce", experiment_name, wandb_api_key)
             self.using_wandb = True
         else:
             self.using_wandb = False

notebooks/enhance_me_train.ipynb

@@ -41,7 +41,8 @@
     "\n",
     "from PIL import Image\n",
     "from enhance_me import commons\n",
-    "from enhance_me.mirnet import MIRNet"
+    "from enhance_me.mirnet import MIRNet\n",
+    "from enhance_me.zero_dce import ZeroDCE"
    ]
   },
   {
@@ -183,7 +184,62 @@
     "id": "dO-IbNQHkB3R"
    },
    "outputs": [],
-   "source": []
+   "source": [
+    "# @title Zero-DCE Train Configs\n",
+    "\n",
+    "experiment_name = \"lol_dataset_128\"  # @param {type:\"string\"}\n",
+    "image_size = 128  # @param {type:\"integer\"}\n",
+    "dataset_label = \"lol\"  # @param [\"lol\"]\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",
+    "use_mixed_precision = False  # @param {type:\"boolean\"}\n",
+    "wandb_api_key = \"\"  # @param {type:\"string\"}\n",
+    "val_split = 0.1  # @param {type:\"slider\", min:0.1, max:1.0, step:0.1}\n",
+    "batch_size = 32  # @param {type:\"integer\"}\n",
+    "learning_rate = 1e-4  # @param {type:\"number\"}\n",
+    "epsilon = 1e-3  # @param {type:\"number\"}\n",
+    "epochs = 100  # @param {type:\"slider\", min:10, max:100, step:5}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "zero_dce = ZeroDCE(\n",
+    "    experiment_name=experiment_name,\n",
+    "    wandb_api_key=None if wandb_api_key == \"\" else wandb_api_key\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "zero_dce.build_datasets(\n",
+    "    image_size=image_size,\n",
+    "    dataset_label=dataset_label,\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",
+    "    val_split=val_split,\n",
+    "    batch_size=batch_size\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "zero_dce.compile(learning_rate=learning_rate)\n",
+    "zero_dce.train(epochs=epochs)"
+   ]
   }
  ],
  "metadata": {