tensorflow_graphics/notebooks/spherical_harmonics_optimization.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ivpj9EeVLWTL"
      },
      "source": [
        "##### Copyright 2019 Google LLC."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "jXIU2YPcN75z"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "UUiCfWD_i6Ck"
      },
      "source": [
        "#Spherical Harmonics Optimization\n",
        "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/graphics/blob/master/tensorflow_graphics/notebooks/spherical_harmonics_optimization.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003eRun in Google Colab\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://github.com/tensorflow/graphics/blob/master/tensorflow_graphics/notebooks/spherical_harmonics_optimization.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003eView source on GitHub\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "\u003c/table\u003e"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "iFqF3lDo3fIa"
      },
      "source": [
        "This Colab covers an advanced topic and hence focuses on providing a toy example\n",
        "to form a high level understanding of how to estimate environment lighting using\n",
        "Spherical Harmonics rather than providing step by step details. We refer the\n",
        "interested reader to\n",
        "[another Colab](https://colab.research.google.com/github/tensorflow/graphics/blob/master/tensorflow_graphics/notebooks/spherical_harmonics_approximation.ipynb)\n",
        "to get a high level understanding of Spherical Harmonics.\n",
        "\n",
        "Given an image of a known object (sphere) with a known reflectance function,\n",
        "this Colab illustrates how to perform optimization of spherical harmonics to\n",
        "recover the lighting environment."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "hM73at_pLadB"
      },
      "source": [
        "## Setup \u0026 Imports\n",
        "If Tensorflow Graphics is not installed on your system, the following cell can install the Tensorflow Graphics package for you."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "G3WpJC9GLfUe"
      },
      "outputs": [],
      "source": [
        "!pip install tensorflow_graphics"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "O7Fe-GqGLif6"
      },
      "source": [
        "Now that Tensorflow Graphics is installed, let's import everything needed to run the demo contained in this notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "cV4dat2Ayjfu"
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "\n",
        "from tensorflow_graphics.rendering.camera import orthographic\n",
        "from tensorflow_graphics.geometry.representation import grid\n",
        "from tensorflow_graphics.geometry.representation import ray\n",
        "from tensorflow_graphics.geometry.representation import vector\n",
        "from tensorflow_graphics.math import spherical_harmonics\n",
        "from tensorflow_graphics.math import math_helpers as tf_math\n",
        "\n",
        "tf.compat.v1.enable_v2_behavior()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Efb6MPaEQQ-m"
      },
      "outputs": [],
      "source": [
        "def compute_intersection_normal_sphere(width, height, radius,\n",
        "                                       center, dtype):\n",
        "  \"\"\"Estimates sphere normal and depth for each pixel in the image.\"\"\"\n",
        "  # Generates a 2d grid storing pixel coordinates.\n",
        "  pixel_grid_start = np.array((0.5, 0.5), dtype=dtype)\n",
        "  pixel_grid_end = np.array((width - 0.5, height - 0.5), dtype=dtype)\n",
        "  pixel_nb = np.array((width, height))\n",
        "  pixels = grid.generate(pixel_grid_start, pixel_grid_end, pixel_nb)\n",
        "  # Computes the ray direction of each pixel.\n",
        "  pixel_ray = tf.math.l2_normalize(orthographic.ray(pixels), axis=-1)\n",
        "  # Defines the position of pixels in world coordinates.\n",
        "  zero_depth = np.zeros([width, height, 1])\n",
        "  pixels_3d = orthographic.unproject(pixels, zero_depth)\n",
        "  # Computes intersections with the sphere and surface normals for each ray.\n",
        "  intersection, normal = ray.intersection_ray_sphere(\n",
        "      center, radius, pixel_ray, pixels_3d)\n",
        "  # Extracts data about the closest intersection.\n",
        "  intersection = intersection[0, ...]\n",
        "  normal = normal[0, ...]\n",
        "  # Replaces NaNs with zeros.\n",
        "  zeros = tf.zeros_like(pixels_3d)\n",
        "  intersection = tf.where(tf.math.is_nan(intersection), zeros, intersection)\n",
        "  normal = tf.where(tf.math.is_nan(normal), zeros, normal)\n",
        "  return intersection, normal"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Mhj6ZBssOvYT"
      },
      "source": [
        "## Spherical Harmonics optimization"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Uvcgy69dzsS1"
      },
      "outputs": [],
      "source": [
        "light_image_width = 100\n",
        "light_image_height = 100\n",
        "dtype = np.float64\n",
        "\n",
        "############################################################################\n",
        "# Builds the pixels grid and computes corresponding spherical coordinates. #\n",
        "############################################################################\n",
        "pixel_grid_start = np.array((0, 0), dtype=dtype)\n",
        "pixel_grid_end = np.array((light_image_width - 1, light_image_height - 1),\n",
        "                          dtype=dtype)\n",
        "pixel_nb = np.array((light_image_width, light_image_height))\n",
        "pixels = grid.generate(pixel_grid_start, pixel_grid_end, pixel_nb)\n",
        "normalized_pixels = pixels / (light_image_width - 1, light_image_height - 1)\n",
        "spherical_coordinates = tf_math.square_to_spherical_coordinates(\n",
        "    normalized_pixels)\n",
        "theta = spherical_coordinates[:, :, 1]\n",
        "phi = spherical_coordinates[:, :, 2]\n",
        "\n",
        "################################################################################################\n",
        "# Builds the Spherical Harmonics and sets coefficients for the light and reflectance functions. #\n",
        "################################################################################################\n",
        "max_band = 2\n",
        "l, m = spherical_harmonics.generate_l_m_permutations(max_band)\n",
        "l_broadcasted = tf.broadcast_to(l, [light_image_width, light_image_height] +\n",
        "                                l.shape.as_list())\n",
        "m_broadcasted = tf.broadcast_to(m, [light_image_width, light_image_height] +\n",
        "                                l.shape.as_list())\n",
        "theta = tf.expand_dims(theta, axis=-1)\n",
        "theta_broadcasted = tf.broadcast_to(theta,\n",
        "                                    [light_image_width, light_image_height, 1])\n",
        "phi = tf.expand_dims(phi, axis=-1)\n",
        "phi_broadcasted = tf.broadcast_to(phi,\n",
        "                                  [light_image_width, light_image_height, 1])\n",
        "sh_coefficients = spherical_harmonics.evaluate_spherical_harmonics(\n",
        "    l_broadcasted, m_broadcasted, theta_broadcasted, phi_broadcasted)\n",
        "\n",
        "# The lighting and BRDF coefficients come from the first Colab demo on Spherical\n",
        "# Harmonics.\n",
        "light_coeffs = np.array((2.17136424e-01, -2.06274278e-01, 3.10378283e-17,\n",
        "                         2.76236879e-01, -3.08694040e-01, -4.69862940e-17,\n",
        "                         -1.85866463e-01, 7.05744675e-17, 9.14290771e-02))\n",
        "brdf_coeffs = np.array((0.28494423, 0.33231551, 0.16889377))\n",
        "\n",
        "# Reconstruction of the light function.\n",
        "reconstructed_light_function = tf.squeeze(\n",
        "    vector.dot(sh_coefficients, light_coeffs))\n",
        "\n",
        "###################################\n",
        "# Setup the image, and the sphere #\n",
        "###################################\n",
        "# Image dimensions\n",
        "image_width = 100\n",
        "image_height = 80\n",
        "\n",
        "# Sphere center and radius\n",
        "sphere_radius = np.array((30,), dtype=dtype)\n",
        "sphere_center = np.array((image_width / 2.0, image_height / 2.0, 100.0),\n",
        "                         dtype=dtype)\n",
        "\n",
        "# Builds the pixels grid and compute corresponding spherical coordinates.\n",
        "pixel_grid_start = np.array((0, 0), dtype=dtype)\n",
        "pixel_grid_end = np.array((image_width - 1, image_height - 1), dtype=dtype)\n",
        "pixel_nb = np.array((image_width, image_height))\n",
        "pixels = grid.generate(pixel_grid_start, pixel_grid_end, pixel_nb)\n",
        "normalized_pixels = pixels / (image_width - 1, image_height - 1)\n",
        "spherical_coordinates = tf_math.square_to_spherical_coordinates(\n",
        "    normalized_pixels)\n",
        "\n",
        "################################################################################################\n",
        "# For each pixel in the image, estimate the corresponding surface point and associated normal. #\n",
        "################################################################################################\n",
        "intersection_3d, surface_normal = compute_intersection_normal_sphere(\n",
        "    image_width, image_height, sphere_radius, sphere_center, dtype)\n",
        "surface_normals_spherical_coordinates = tf_math.cartesian_to_spherical_coordinates(\n",
        "    surface_normal)\n",
        "\n",
        "##########################################\n",
        "# Estimates result using SH convolution. #\n",
        "##########################################\n",
        "target = spherical_harmonics.integration_product(\n",
        "    light_coeffs,\n",
        "    spherical_harmonics.rotate_zonal_harmonics(\n",
        "        brdf_coeffs,\n",
        "        tf.expand_dims(surface_normals_spherical_coordinates[:, :, 1], axis=-1),\n",
        "        tf.expand_dims(surface_normals_spherical_coordinates[:, :, 2],\n",
        "                       axis=-1)),\n",
        "    keepdims=False)\n",
        "# Sets pixels not belonging to the sphere to 0.\n",
        "target = tf.where(\n",
        "    tf.greater(intersection_3d[:, :, 2], 0.0), target, tf.zeros_like(target))\n",
        "\n",
        "#########################################################################################\n",
        "# Optimization of the lighting coefficients by minimization of the reconstruction error #\n",
        "#########################################################################################\n",
        "# Initial solution.\n",
        "recovered_light_coeffs = tf.Variable(\n",
        "    np.array((1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0)))\n",
        "\n",
        "\n",
        "def reconstruct_image(recovered_light_coeffs):\n",
        "  reconstructed_image = spherical_harmonics.integration_product(\n",
        "      recovered_light_coeffs,\n",
        "      spherical_harmonics.rotate_zonal_harmonics(\n",
        "          brdf_coeffs,\n",
        "          tf.expand_dims(\n",
        "              surface_normals_spherical_coordinates[:, :, 1], axis=-1),\n",
        "          tf.expand_dims(\n",
        "              surface_normals_spherical_coordinates[:, :, 2], axis=-1)),\n",
        "      keepdims=False)\n",
        "  return tf.where(\n",
        "      tf.greater(intersection_3d[:, :, 2], 0.0), reconstructed_image,\n",
        "      tf.zeros_like(target))\n",
        "\n",
        "\n",
        "# Sets the optimization problem up.\n",
        "def my_loss(recovered_light_coeffs):\n",
        "  reconstructed_image = reconstruct_image(recovered_light_coeffs)\n",
        "  return tf.nn.l2_loss(reconstructed_image - target) / (\n",
        "      image_width * image_height)\n",
        "\n",
        "\n",
        "learning_rate = 0.1\n",
        "with tf.name_scope(\"optimization\"):\n",
        "  optimizer = tf.compat.v1.train.AdamOptimizer(learning_rate)\n",
        "\n",
        "\n",
        "def gradient_loss(recovered_light_coeffs):\n",
        "  with tf.GradientTape() as tape:\n",
        "    loss_value = my_loss(recovered_light_coeffs)\n",
        "  return tape.gradient(loss_value, [recovered_light_coeffs])\n",
        "\n",
        "\n",
        "####################\n",
        "# Initial solution #\n",
        "####################\n",
        "target_transpose = np.transpose(target, (1, 0))\n",
        "reconstructed_image = reconstruct_image(recovered_light_coeffs)\n",
        "reconstructed_image = np.transpose(reconstructed_image, (1, 0))\n",
        "plt.figure(figsize=(10, 20))\n",
        "ax = plt.subplot(1, 3, 1)\n",
        "ax.axes.get_xaxis().set_visible(False)\n",
        "ax.axes.get_yaxis().set_visible(False)\n",
        "ax.grid(False)\n",
        "ax.set_title(\"Target\")\n",
        "_ = ax.imshow(target_transpose, vmin=0.0)\n",
        "ax = plt.subplot(1, 3, 2)\n",
        "ax.axes.get_xaxis().set_visible(False)\n",
        "ax.axes.get_yaxis().set_visible(False)\n",
        "ax.grid(False)\n",
        "ax.set_title(\"Initial solution\")\n",
        "_ = ax.imshow(reconstructed_image, vmin=0.0)\n",
        "\n",
        "################\n",
        "# Optimization #\n",
        "################\n",
        "nb_iterations = 100\n",
        "for it in range(nb_iterations):\n",
        "  gradients_loss = gradient_loss(recovered_light_coeffs)\n",
        "  optimizer.apply_gradients(zip(gradients_loss, (recovered_light_coeffs,)))\n",
        "\n",
        "  if it % 33 == 0:\n",
        "    reconstructed_image = reconstruct_image(recovered_light_coeffs)\n",
        "    reconstructed_image = np.transpose(reconstructed_image, (1, 0))\n",
        "    # Displays the target and prediction.\n",
        "    plt.figure(figsize=(10, 20))\n",
        "    ax = plt.subplot(1, 3, 1)\n",
        "    ax.axes.get_xaxis().set_visible(False)\n",
        "    ax.axes.get_yaxis().set_visible(False)\n",
        "    ax.grid(False)\n",
        "    ax.set_title(\"Target\")\n",
        "    img = ax.imshow(target_transpose, vmin=0.0)\n",
        "    ax = plt.subplot(1, 3, 2)\n",
        "    ax.axes.get_xaxis().set_visible(False)\n",
        "    ax.axes.get_yaxis().set_visible(False)\n",
        "    ax.grid(False)\n",
        "    ax.set_title(\"Prediction iteration \" + str(it))\n",
        "    img = ax.imshow(reconstructed_image, vmin=0.0)\n",
        "    # Shows the difference between groundtruth and prediction.\n",
        "    vmax = np.maximum(np.amax(reconstructed_image), np.amax(target_transpose))\n",
        "    ax = plt.subplot(1, 3, 3)\n",
        "    ax.axes.get_xaxis().set_visible(False)\n",
        "    ax.axes.get_yaxis().set_visible(False)\n",
        "    ax.grid(False)\n",
        "    ax.set_title(\"Difference iteration \" + str(it))\n",
        "    img = ax.imshow(\n",
        "        np.abs(reconstructed_image - target_transpose), vmin=0.0, vmax=vmax)\n",
        "\n",
        "# Reconstructs the groundtruth and predicted environment maps.\n",
        "reconstructed_predicted_light = tf.squeeze(\n",
        "    vector.dot(sh_coefficients, recovered_light_coeffs))\n",
        "\n",
        "# Displays the groundtruth and predicted environment maps.\n",
        "plt.figure(figsize=(10, 20))\n",
        "ax = plt.subplot(1, 2, 1)\n",
        "ax.axes.get_xaxis().set_visible(False)\n",
        "ax.axes.get_yaxis().set_visible(False)\n",
        "ax.grid(False)\n",
        "ax.set_title(\"Target light\")\n",
        "img = ax.imshow(reconstructed_light_function, vmin=0.0)\n",
        "ax = plt.subplot(1, 2, 2)\n",
        "ax.axes.get_xaxis().set_visible(False)\n",
        "ax.axes.get_yaxis().set_visible(False)\n",
        "ax.grid(False)\n",
        "ax.set_title(\"Predicted light\")\n",
        "img = ax.imshow(reconstructed_predicted_light, vmin=0.0)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "last_runtime": {
        "build_target": "//learning/brain/python/client:colab_notebook",
        "kind": "private"
      },
      "name": "spherical harmonics optimization.ipynb",
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}