updated notebook

notebooks/enhance_me_train.ipynb

@@ -22,7 +22,7 @@
    },
    "outputs": [],
    "source": [
-    "!git clone https://github.com/soumik12345/enhance-me\n",
+    "!git clone https://github.com/soumik12345/enhance-me -b zero-dce\n",
     "!pip install -qqq wandb streamlit"
    ]
   },
@@ -171,7 +171,7 @@
     "    enhanced_image = mirnet.infer(original_image)\n",
     "    ground_truth = Image.open(mirnet.test_enhanced_images[index])\n",
     "    commons.plot_results(\n",
-    "        [original_image, ground_truth, ground_truth],\n",
+    "        [original_image, ground_truth, enhanced_image],\n",
     "        [\"Original Image\", \"Ground Truth\", \"Enhanced Image\"],\n",
     "        (18, 18),\n",
     "    )"
@@ -238,7 +238,24 @@
    "outputs": [],
    "source": [
     "zero_dce.compile(learning_rate=learning_rate)\n",
-    "zero_dce.train(epochs=epochs)"
+    "history = zero_dce.train(epochs=epochs)\n",
+    "zero_dce.save_weights(os.path.join(experiment_name, \"weights.h5\"))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "for index, low_image_file in enumerate(zero_dce.test_low_images):\n",
+    "    original_image = Image.open(low_image_file)\n",
+    "    enhanced_image = zero_dce.infer(original_image)\n",
+    "    commons.plot_results(\n",
+    "        [original_image, enhanced_image],\n",
+    "        [\"Original Image\", \"Enhanced Image\"],\n",
+    "        (18, 18),\n",
+    "    )"
    ]
   }
  ],