recoloring.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "authorship_tag": "ABX9TyPD3pXRpOuigYiFeN4QIF32",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<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": "code",
      "execution_count": null,
      "metadata": {
        "id": "nnOA8b8E3Y6U"
      },
      "outputs": [],
      "source": [
        "!git clone https://github.com/compphoto/Intrinsic\n",
        "!cd Intrinsic/ && pip install ."
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "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, 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": "mS9AFGEj3jhc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "intrinsic_model = load_models('paper_weights')"
      ],
      "metadata": {
        "id": "nw0poq363mqy"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# three different example scenes from the paper\n",
        "scene_name = 'yellow_chair'\n",
        "# scene_name = 'brown_chairs'\n",
        "# scene_name = 'spain_museum'"
      ],
      "metadata": {
        "id": "l5YBpZ3a5rfS"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "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": "zmWeGYSP77W4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "results = run_pipeline(\n",
        "    intrinsic_model,\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": "pDr4Wh7M5pIE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "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": "kip8Y1tE83Zi"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# 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": "oF2ljuOv84WA"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}