recoloring.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "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",
   "execution_count": null,
   "metadata": {
    "id": "mS9AFGEj3jhc"
   },
   "outputs": [],
   "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_gray_pipeline, load_models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "nw0poq363mqy"
   },
   "outputs": [],
   "source": [
    "intrinsic_model = load_models('paper_weights')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "l5YBpZ3a5rfS"
   },
   "outputs": [],
   "source": [
    "# three different example scenes from the paper\n",
    "scene_name = 'yellow_chair'\n",
    "# scene_name = 'brown_chairs'\n",
    "# scene_name = 'spain_museum'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "zmWeGYSP77W4"
   },
   "outputs": [],
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "pDr4Wh7M5pIE"
   },
   "outputs": [],
   "source": [
    "results = run_gray_pipeline(\n",
    "    intrinsic_model,\n",
    "    inp,\n",
    "    resize_conf=None,\n",
    "    maintain_size=True\n",
    ")\n",
    "\n",
    "alb = results['gry_alb']\n",
    "image = results['image']\n",
    "inv_shd = results['gry_shd']\n",
    "\n",
    "shd = uninvert(inv_shd)[:, :, None]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "kip8Y1tE83Zi"
   },
   "outputs": [],
   "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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "oF2ljuOv84WA"
   },
   "outputs": [],
   "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": {
  "accelerator": "GPU",
  "colab": {
   "authorship_tag": "ABX9TyPD3pXRpOuigYiFeN4QIF32",
   "gpuType": "T4",
   "include_colab_link": true,
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}