updated mirnet

enhance_me/commons.py

@@ -43,7 +43,7 @@ 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
-        wandb.init(project=project_name, name=experiment_name)
+        wandb.init(project=project_name, name=experiment_name, sync_tensorboard=True)
 
 
 def download_lol_dataset():

enhance_me/mirnet/mirnet.py

@@ -21,11 +21,7 @@ from ..commons import (
 
 
 class MIRNet:
-    def __init__(
-        self,
-        experiment_name: str,
-        wandb_api_key=None,
-    ) -> None:
+    def __init__(self, experiment_name: str, wandb_api_key=None) -> None:
         self.experiment_name = experiment_name
         if wandb_api_key is not None:
             init_wandb("mirnet", experiment_name, wandb_api_key)

notebooks/enhance_me_train.ipynb

@@ -1,284 +1,197 @@
 {
-  "nbformat": 4,
-  "nbformat_minor": 0,
-  "metadata": {
-    "accelerator": "GPU",
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "colab_type": "text",
+    "id": "view-in-github"
+   },
+   "source": [
+    "<a href=\"https://colab.research.google.com/github/soumik12345/enhance-me/blob/mirnet/notebooks/enhance_me_train.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
     "colab": {
-      "name": "enhance-me-train.ipynb",
-      "provenance": [],
-      "collapsed_sections": [],
-      "machine_shape": "hm",
-      "authorship_tag": "ABX9TyN4LuJh6kWhbqxzA5s9sp7k",
-      "include_colab_link": true
-    },
-    "kernelspec": {
-      "display_name": "Python 3",
-      "name": "python3"
+     "base_uri": "https://localhost:8080/"
     },
-    "language_info": {
-      "name": "python"
-    }
+    "id": "1JryaVhtBHij",
+    "outputId": "97ee6a4a-2479-4124-e96a-f0a792bdec46"
+   },
+   "outputs": [],
+   "source": [
+    "!git clone https://github.com/soumik12345/enhance-me -b mirnet\n",
+    "!pip install -qqq wandb streamlit"
+   ]
   },
-  "cells": [
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "view-in-github",
-        "colab_type": "text"
-      },
-      "source": [
-        "<a href=\"https://colab.research.google.com/github/soumik12345/enhance-me/blob/mirnet/notebooks/enhance_me_train.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/"
-        },
-        "id": "1JryaVhtBHij",
-        "outputId": "97ee6a4a-2479-4124-e96a-f0a792bdec46"
-      },
-      "source": [
-        "!git clone https://github.com/soumik12345/enhance-me -b mirnet\n",
-        "!pip install -qqq wandb streamlit"
-      ],
-      "execution_count": 1,
-      "outputs": [
-        {
-          "output_type": "stream",
-          "name": "stdout",
-          "text": [
-            "Cloning into 'enhance-me'...\n",
-            "remote: Enumerating objects: 89, done.\u001b[K\n",
-            "remote: Counting objects: 100% (89/89), done.\u001b[K\n",
-            "remote: Compressing objects: 100% (61/61), done.\u001b[K\n",
-            "remote: Total 89 (delta 43), reused 63 (delta 23), pack-reused 0\u001b[K\n",
-            "Unpacking objects: 100% (89/89), done.\n",
-            "\u001b[K     |████████████████████████████████| 1.7 MB 8.2 MB/s \n",
-            "\u001b[K     |████████████████████████████████| 9.1 MB 33.4 MB/s \n",
-            "\u001b[K     |████████████████████████████████| 140 kB 74.7 MB/s \n",
-            "\u001b[K     |████████████████████████████████| 97 kB 8.6 MB/s \n",
-            "\u001b[K     |████████████████████████████████| 180 kB 83.6 MB/s \n",
-            "\u001b[K     |████████████████████████████████| 63 kB 2.1 MB/s \n",
-            "\u001b[K     |████████████████████████████████| 4.3 MB 83.4 MB/s \n",
-            "\u001b[K     |████████████████████████████████| 178 kB 68.0 MB/s \n",
-            "\u001b[K     |████████████████████████████████| 76 kB 6.0 MB/s \n",
-            "\u001b[K     |████████████████████████████████| 111 kB 81.8 MB/s \n",
-            "\u001b[K     |████████████████████████████████| 125 kB 86.7 MB/s \n",
-            "\u001b[K     |████████████████████████████████| 791 kB 67.2 MB/s \n",
-            "\u001b[K     |████████████████████████████████| 374 kB 83.4 MB/s \n",
-            "\u001b[?25h  Building wheel for subprocess32 (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
-            "  Building wheel for pympler (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
-            "  Building wheel for blinker (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
-            "  Building wheel for pathtools (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
-            "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
-            "jupyter-console 5.2.0 requires prompt-toolkit<2.0.0,>=1.0.0, but you have prompt-toolkit 3.0.23 which is incompatible.\n",
-            "google-colab 1.0.0 requires ipykernel~=4.10, but you have ipykernel 6.5.1 which is incompatible.\n",
-            "google-colab 1.0.0 requires ipython~=5.5.0, but you have ipython 7.30.0 which is incompatible.\u001b[0m\n"
-          ]
-        }
-      ]
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "G_c4VtXWHR5l"
-      },
-      "source": [
-        "import sys\n",
-        "sys.path.append(\"./enhance-me\")\n",
-        "\n",
-        "from PIL import Image\n",
-        "from enhance_me import commons\n",
-        "from enhance_me.mirnet import MIRNet"
-      ],
-      "execution_count": 2,
-      "outputs": []
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "ZpBHbYaMIqP_"
-      },
-      "source": [
-        "#@title MIRNet Train Configs\n",
-        "\n",
-        "experiment_name = 'lol_dataset_256' #@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",
-        "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 = 4 #@param {type:\"integer\"}\n",
-        "num_recursive_residual_groups = 3 #@param {type:\"slider\", min:1, max:5, step:1}\n",
-        "num_multi_scale_residual_blocks = 2 #@param {type:\"slider\", min:1, max:5, step:1}\n",
-        "learning_rate = 1e-4 #@param {type:\"number\"}\n",
-        "epsilon = 1e-3 #@param {type:\"number\"}\n",
-        "epochs = 50 #@param {type:\"slider\", min:10, max:100, step:5}"
-      ],
-      "execution_count": 3,
-      "outputs": []
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/",
-          "height": 52
-        },
-        "id": "IVRoedqBIMuH",
-        "outputId": "53ca5beb-871a-4ec3-b757-173e09a15331"
-      },
-      "source": [
-        "mirnet = MIRNet(\n",
-        "    experiment_name=experiment_name,\n",
-        "    wandb_api_key=None if wandb_api_key == '' else wandb_api_key\n",
-        ")"
-      ],
-      "execution_count": 4,
-      "outputs": [
-        {
-          "output_type": "stream",
-          "name": "stderr",
-          "text": [
-            "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33m19soumik-rakshit96\u001b[0m (use `wandb login --relogin` to force relogin)\n"
-          ]
-        },
-        {
-          "output_type": "display_data",
-          "data": {
-            "text/html": [
-              "\n",
-              "                    Syncing run <strong><a href=\"https://wandb.ai/19soumik-rakshit96/mirnet/runs/3p3rc341\" target=\"_blank\">lol_dataset_256</a></strong> to <a href=\"https://wandb.ai/19soumik-rakshit96/mirnet\" target=\"_blank\">Weights & Biases</a> (<a href=\"https://docs.wandb.com/integrations/jupyter.html\" target=\"_blank\">docs</a>).<br/>\n",
-              "\n",
-              "                "
-            ],
-            "text/plain": [
-              "<IPython.core.display.HTML object>"
-            ]
-          },
-          "metadata": {}
-        }
-      ]
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/"
-        },
-        "id": "O66Iwzx8IsGh",
-        "outputId": "0b6f1683-65d1-4737-a32f-d36b331d2bc2"
-      },
-      "source": [
-        "mirnet.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",
-        ")"
-      ],
-      "execution_count": 5,
-      "outputs": [
-        {
-          "output_type": "stream",
-          "name": "stdout",
-          "text": [
-            "Downloading data from https://github.com/soumik12345/enhance-me/releases/download/v0.1/lol_dataset.zip\n",
-            "347176960/347171015 [==============================] - 13s 0us/step\n",
-            "347185152/347171015 [==============================] - 13s 0us/step\n",
-            "Number of train data points: 436\n",
-            "Number of validation data points: 49\n"
-          ]
-        }
-      ]
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "G_c4VtXWHR5l"
+   },
+   "outputs": [],
+   "source": [
+    "import sys\n",
+    "\n",
+    "sys.path.append(\"./enhance-me\")\n",
+    "\n",
+    "from PIL import Image\n",
+    "from enhance_me import commons\n",
+    "from enhance_me.mirnet import MIRNet"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "ZpBHbYaMIqP_"
+   },
+   "outputs": [],
+   "source": [
+    "# @title MIRNet Train Configs\n",
+    "\n",
+    "experiment_name = \"lol_dataset_256\"  # @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",
+    "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 = 4  # @param {type:\"integer\"}\n",
+    "num_recursive_residual_groups = 3  # @param {type:\"slider\", min:1, max:5, step:1}\n",
+    "num_multi_scale_residual_blocks = 2  # @param {type:\"slider\", min:1, max:5, step:1}\n",
+    "learning_rate = 1e-4  # @param {type:\"number\"}\n",
+    "epsilon = 1e-3  # @param {type:\"number\"}\n",
+    "epochs = 50  # @param {type:\"slider\", min:10, max:100, step:5}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/",
+     "height": 52
     },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "tsfKrBCsL_Bb"
-      },
-      "source": [
-        "mirnet.build_model(\n",
-        "    num_recursive_residual_groups=num_recursive_residual_groups,\n",
-        "    num_multi_scale_residual_blocks=num_multi_scale_residual_blocks,\n",
-        "    learning_rate=learning_rate,\n",
-        "    epsilon=epsilon\n",
-        ")"
-      ],
-      "execution_count": 6,
-      "outputs": []
+    "id": "IVRoedqBIMuH",
+    "outputId": "53ca5beb-871a-4ec3-b757-173e09a15331"
+   },
+   "outputs": [],
+   "source": [
+    "mirnet = MIRNet(\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": {
+    "colab": {
+     "base_uri": "https://localhost:8080/"
     },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/"
-        },
-        "id": "y3L9wlpkNziL",
-        "outputId": "5149f0e7-91f4-450f-c43a-1b6028692bbc"
-      },
-      "source": [
-        "history = mirnet.train(epochs=epochs)"
-      ],
-      "execution_count": null,
-      "outputs": [
-        {
-          "output_type": "stream",
-          "name": "stderr",
-          "text": [
-            "/usr/local/lib/python3.7/dist-packages/keras/engine/functional.py:1410: CustomMaskWarning: Custom mask layers require a config and must override get_config. When loading, the custom mask layer must be passed to the custom_objects argument.\n",
-            "  layer_config = serialize_layer_fn(layer)\n"
-          ]
-        },
-        {
-          "output_type": "stream",
-          "name": "stdout",
-          "text": [
-            "Epoch 1/50\n",
-            " 66/218 [========>.....................] - ETA: 2:25 - loss: 0.1721 - peak_signal_noise_ratio: 63.2555"
-          ]
-        }
-      ]
+    "id": "O66Iwzx8IsGh",
+    "outputId": "0b6f1683-65d1-4737-a32f-d36b331d2bc2"
+   },
+   "outputs": [],
+   "source": [
+    "mirnet.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": {
+    "id": "tsfKrBCsL_Bb"
+   },
+   "outputs": [],
+   "source": [
+    "mirnet.build_model(\n",
+    "    num_recursive_residual_groups=num_recursive_residual_groups,\n",
+    "    num_multi_scale_residual_blocks=num_multi_scale_residual_blocks,\n",
+    "    learning_rate=learning_rate,\n",
+    "    epsilon=epsilon,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/"
     },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "colab": {
-          "background_save": true
-        },
-        "id": "daFKbgBkiyzc"
-      },
-      "source": [
-        "for index, low_image_file in enumerate(mirnet.test_low_images):\n",
-        "    original_image = Image.open(low_image_file)\n",
-        "    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",
-        "        (18, 18)\n",
-        "    )"
-      ],
-      "execution_count": null,
-      "outputs": []
+    "id": "y3L9wlpkNziL",
+    "outputId": "5149f0e7-91f4-450f-c43a-1b6028692bbc"
+   },
+   "outputs": [],
+   "source": [
+    "history = mirnet.train(epochs=epochs)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "colab": {
+     "background_save": true
     },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "dO-IbNQHkB3R"
-      },
-      "source": [
-        ""
-      ],
-      "execution_count": null,
-      "outputs": []
-    }
-  ]
-}
+    "id": "daFKbgBkiyzc"
+   },
+   "outputs": [],
+   "source": [
+    "for index, low_image_file in enumerate(mirnet.test_low_images):\n",
+    "    original_image = Image.open(low_image_file)\n",
+    "    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",
+    "        (18, 18),\n",
+    "    )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "dO-IbNQHkB3R"
+   },
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "accelerator": "GPU",
+  "colab": {
+   "authorship_tag": "ABX9TyN4LuJh6kWhbqxzA5s9sp7k",
+   "collapsed_sections": [],
+   "include_colab_link": true,
+   "machine_shape": "hm",
+   "name": "enhance-me-train.ipynb",
+   "provenance": []
+  },
+  "kernelspec": {
+   "display_name": "Python 3",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
+}