notebook for recoloring

Author: CCareaga

recoloring.ipynb

@@ -5,6 +5,7 @@
     "colab": {
       "provenance": [],
       "gpuType": "T4",
+      "authorship_tag": "ABX9TyPD3pXRpOuigYiFeN4QIF32",
       "include_colab_link": true
     },
     "kernelspec": {
@@ -27,27 +28,16 @@
         "<a href=\"https://colab.research.google.com/github/compphoto/Intrinsic/blob/main/recoloring.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
       ]
     },
-    {
-      "cell_type": "markdown",
-      "source": [
-        "Make sure to set the runtime to GPU: Runtime -> Change runtime type -> T4 GPU\n",
-        "\n",
-        "You can upload your own images, then change the relevant code cells to load it and send it through the model."
-      ],
-      "metadata": {
-        "id": "-lKPW1wZEulh"
-      }
-    },
     {
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "_cPkUxN-ikld"
+        "id": "nnOA8b8E3Y6U"
       },
       "outputs": [],
       "source": [
-        "# install the intrinsic decomposition repo from github\n",
-        "!pip install https://github.com/compphoto/Intrinsic/archive/main.zip"
+        "!git clone https://github.com/compphoto/Intrinsic\n",
+        "!cd Intrinsic/ && pip install ."
       ]
     },
     {
@@ -56,95 +46,122 @@
         "import torch\n",
         "\n",
         "# import some helper functions from chrislib (will be installed by the intrinsic repo)\n",
-        "from chrislib.general import show, view, uninvert\n",
+        "from chrislib.general import show, view, uninvert, match_scale\n",
         "from chrislib.data_util import load_image\n",
         "\n",
         "# import model loading and running the pipeline\n",
         "from intrinsic.pipeline import run_pipeline\n",
         "from intrinsic.model_util import load_models"
       ],
       "metadata": {
-        "id": "J0gn82ZSjomn"
+        "id": "mS9AFGEj3jhc"
       },
       "execution_count": null,
       "outputs": []
     },
     {
       "cell_type": "code",
       "source": [
-        "# download the pretrained weights and return the model (may take a bit to download weights)\n",
         "intrinsic_model = load_models('paper_weights')"
       ],
       "metadata": {
-        "id": "Ap3HubpwC_KG"
+        "id": "nw0poq363mqy"
       },
       "execution_count": null,
       "outputs": []
     },
     {
       "cell_type": "code",
       "source": [
-        "# load an example image from the github repo\n",
-        "torch.hub.download_url_to_file('https://raw.githubusercontent.com/compphoto/Intrinsic/main/figures/avocado.png', 'avo.png')"
+        "# three different example scenes from the paper\n",
+        "scene_name = 'yellow_chair'\n",
+        "# scene_name = 'brown_chairs'\n",
+        "# scene_name = 'spain_museum'"
       ],
       "metadata": {
-        "id": "m_NYfDx0AhTw"
+        "id": "l5YBpZ3a5rfS"
       },
       "execution_count": null,
       "outputs": []
     },
     {
       "cell_type": "code",
       "source": [
-        "# load the image to run through the pipeline\n",
-        "img = load_image('/content/avo.png')"
+        "inp = load_image(f'Intrinsic/examples/{scene_name}/input.png')[:, :, :3]\n",
+        "msk = load_image(f'Intrinsic/examples/{scene_name}/mask.png')[:, :, :3]\n",
+        "tex = load_image(f'Intrinsic/examples/{scene_name}/texture.png')[:, :, :3] ** 2.2"
       ],
       "metadata": {
-        "id": "ALb4Pjfvj-MU"
+        "id": "zmWeGYSP77W4"
       },
       "execution_count": null,
       "outputs": []
     },
     {
       "cell_type": "code",
       "source": [
-        "# run the image through the pipeline (use R0 resizing dicussed in the paper)\n",
-        "result = run_pipeline(\n",
+        "results = run_pipeline(\n",
         "    intrinsic_model,\n",
-        "    img,\n",
-        "    resize_conf=0.0,\n",
-        "    maintain_size=True,\n",
-        "    linear=False,\n",
-        "    device='cuda'\n",
-        ")"
+        "    inp,\n",
+        "    resize_conf=None,\n",
+        "    maintain_size=True\n",
+        ")\n",
+        "\n",
+        "alb = results['albedo']\n",
+        "image = results['image']\n",
+        "inv_shd = results['inv_shading']\n",
+        "\n",
+        "shd = uninvert(inv_shd)[:, :, None]"
       ],
       "metadata": {
-        "id": "QW0TiFypkOj-"
+        "id": "pDr4Wh7M5pIE"
       },
       "execution_count": null,
       "outputs": []
     },
     {
       "cell_type": "code",
       "source": [
-        "# convert the inverse shading to regular shading for visualization\n",
-        "shd = uninvert(result['inv_shading'])\n",
-        "alb = result['albedo']"
+        "def perform_recolor(msk, alb, shd, shd_power=1.0, recolor=None):\n",
+        "    # this function will perform the illumination-aware recoloring, or apply a shading curve\n",
+        "    # msk - numpy array (HxWx1) denoting the region to perform the edit\n",
+        "    # alb - linear albedo of the image\n",
+        "    # shd - linear shading of the image\n",
+        "    # shd_power - exponent to apply to the shading (<1 for more diffuse, >1 for more specular)\n",
+        "    # recolor - a texture to apply to the edited region, no recoloring is performed if set to None\n",
+        "\n",
+        "    if recolor is None:\n",
+        "        our_new_alb = alb\n",
+        "    else:\n",
+        "        # we match the scale of the texture to the albedo in the edited region to\n",
+        "        # ensure the appearance of the region is maintained, but this can be altered\n",
+        "        recolor = match_scale(recolor, alb, msk.astype(bool))\n",
+        "        our_new_alb = ((1.0 - msk) * alb) + (msk * recolor)\n",
+        "\n",
+        "    # apply exponentiation to the shading of the region and composite\n",
+        "    masked_shd = msk * (shd ** shd_power)\n",
+        "    new_shd = ((1.0 - msk) * shd) + masked_shd\n",
+        "\n",
+        "    # combine edited albedo and shading, gamma correct and clip\n",
+        "    recolored = (our_new_alb * new_shd) ** (1/2.2)\n",
+        "\n",
+        "    return recolored.clip(0, 1)"
       ],
       "metadata": {
-        "id": "XpYY2MNjkp2f"
+        "id": "kip8Y1tE83Zi"
       },
       "execution_count": null,
       "outputs": []
     },
     {
       "cell_type": "code",
       "source": [
-        "# show the result (gamma corrects the linear intrinsic components and scales to [0-1])\n",
-        "show([img, view(shd), view(alb)], size=(20, 7))"
+        "# NOTE: setting the shading exponent to >1 will make the shading appear more specular,\n",
+        "# but small errors in the shading (albedo leakage) will be amplified in some cases\n",
+        "show(perform_recolor(msk, alb, shd, 1.0, recolor=tex))"
       ],
       "metadata": {
-        "id": "8KKbyoVLki9s"
+        "id": "oF2ljuOv84WA"
       },
       "execution_count": null,
       "outputs": []