3D Projective Geometry and Structured Light_pset6.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/richasempire/AECTECH/blob/main/3D%20Projective%20Geometry%20and%20Structured%20Light_pset6.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Q4dFIVuBrUch"
      },
      "source": [
        "# Problem Set 6: 3D Projective Geometry and Structured Light\n",
        "\n",
        "This assignment is entirely contained within this Jupyter notebook, which you can host on Google Colab. If you're more comfortable with your local dev environment, feel free to download this notebook instead (\"File -> Download -> Download .ipynb\") and edit it in your local environment. Otherwise log into your Google Account and click \"File -> Save a copy in Drive\". You can then edit your copy of the notebook to complete your assignment. You'll find instructions for submitting your work at the [end of this notebook](#scrollTo=Submission_Instructions).\n",
        "\n",
        "This assignment is the same for 6.8300 and 6.8301. Code completion tools are permitted. Your grade on this problem set will be determined by an autograder that runs your code against a set of hidden test cases. Read section 0.3 for details."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "crADUh8p1QHJ"
      },
      "source": [
        "# 0. Setup\n",
        "The next couple of cells import/download the required packages and setup some common methods that you will need for the rest of this problem set."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SZPPfdQ11TwP"
      },
      "source": [
        "## 0.1. Imports"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gwg1s-Qaarre",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "cc45bda7-81f6-40e8-a557-1bab99f28e3d"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Collecting openmesh\n",
            "  Downloading openmesh-1.2.1.tar.gz (9.3 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m9.3/9.3 MB\u001b[0m \u001b[31m18.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.10/dist-packages (from openmesh) (1.25.2)\n",
            "Building wheels for collected packages: openmesh\n",
            "  Building wheel for openmesh (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for openmesh: filename=openmesh-1.2.1-cp310-cp310-linux_x86_64.whl size=855149 sha256=637272deb7ffa1493d5dc337583108c361725fb1e8d59d2c5cb799581e55536d\n",
            "  Stored in directory: /root/.cache/pip/wheels/05/52/ac/4cf307e2dac381ab093a0390901305aee1110ca733ebf55d55\n",
            "Successfully built openmesh\n",
            "Installing collected packages: openmesh\n",
            "Successfully installed openmesh-1.2.1\n"
          ]
        }
      ],
      "source": [
        "# Install openmesh\n",
        "!pip install openmesh\n",
        "\n",
        "# Visualization libraries\n",
        "import tqdm\n",
        "import matplotlib\n",
        "import matplotlib.pyplot as plt\n",
        "from matplotlib import animation, rc\n",
        "\n",
        "%matplotlib inline\n",
        "%config InlineBackend.figure_format='retina'\n",
        "\n",
        "import numpy as np\n",
        "import torch\n",
        "\n",
        "# To make code more readable: some helpful Python 3 type hints\n",
        "from typing import Callable, List, Optional, Tuple, Any, Iterable\n",
        "\n",
        "# We'll use a namedtuple to store scene information in an accessible way\n",
        "from collections import namedtuple\n",
        "\n",
        "# For reproducibility\n",
        "torch.manual_seed(1234)\n",
        "np.random.seed(1234)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0xSnhuim2VZt"
      },
      "source": [
        "## 0.2. Utilities"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "U3b5twqA1HoV"
      },
      "outputs": [],
      "source": [
        "# This pset does not require a CUDA device. It can be done using only a CPU\n",
        "#\n",
        "device = torch.device(\"cpu\")\n",
        "\n",
        "# We'll create a simple namedtuple to hold information about our scene for\n",
        "# Problem 2 (Structured Light)\n",
        "#\n",
        "Scene = namedtuple('Scene',\n",
        "                   ['rgb_ground_truth',\n",
        "                    'z_proj_ground_truth',\n",
        "                    'z_cam_ground_truth',\n",
        "                    'cam_extrinsics',\n",
        "                    'cam_intrinsics',\n",
        "                    'proj_extrinsics',\n",
        "                    'proj_intrinsics'])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "X_khGSBNF9H5"
      },
      "outputs": [],
      "source": [
        "# Setup display utilities\n",
        "def imshow(img, ax=plt, cmap=None):\n",
        "    if len(img.shape) == 3:\n",
        "        # display color image\n",
        "        img_RGB = img[:,:,::-1] # convert BGR to RGB\n",
        "        img_RGB[np.isnan(img_RGB)] = img_RGB[~np.isnan(img_RGB)].min() # visualize nan as the minimal color\n",
        "        axis_image = ax.imshow(img_RGB, cmap=cmap)\n",
        "    else:\n",
        "        # display depth image\n",
        "        img_depth = img.copy()\n",
        "        img_depth[np.isnan(img_depth)] = img_depth[~np.isnan(img_depth)].min() # visualize nan as the minimal depth\n",
        "        axis_image = ax.imshow(img_depth, cmap=cmap)\n",
        "\n",
        "    return axis_image\n",
        "\n",
        "# Animation wrapper to make it easy to display a series of frames in colab\n",
        "class AnimationWrapper():\n",
        "    def __init__(self, rows, cols, frames, figsize=(8, 6)):\n",
        "        self.rows, self.cols = rows, cols\n",
        "        self.frames = frames\n",
        "\n",
        "        self.fig = plt.figure(figsize=figsize)\n",
        "        self.axes = []\n",
        "        self.axis_images = []\n",
        "        self.data = []\n",
        "        self.pbar = None\n",
        "        for r in range(rows):\n",
        "            for c in range(cols):\n",
        "                ax = plt.subplot(rows, cols, r * cols + c + 1)\n",
        "                img = imshow(frames[0][r * cols + c], ax=ax)\n",
        "\n",
        "                self.axes.append(ax)\n",
        "                self.axis_images.append(img)\n",
        "                self.data.append([])\n",
        "        plt.close()\n",
        "\n",
        "        # Set maximal limit to 100 MB, default is 20 MB\n",
        "        matplotlib.rcParams['animation.embed_limit'] = 100\n",
        "\n",
        "    def get_axes(self):\n",
        "        return self.axes\n",
        "\n",
        "    def animate(self, i):\n",
        "        self.pbar.update(1)\n",
        "\n",
        "        for ax_id in range(len(self.axis_images)):\n",
        "            frame = self.frames[i][ax_id]\n",
        "            if len(frame.shape) == 3:\n",
        "                self.axis_images[ax_id].set_data(frame[:,:,::-1])\n",
        "            else:\n",
        "                depth = frame.copy()\n",
        "                depth[np.isnan(depth)] = depth[~np.isnan(depth)].min()\n",
        "                self.axis_images[ax_id].set_data(depth)\n",
        "                self.axis_images[ax_id].set_clim(depth.min(), depth.max())\n",
        "        return tuple(self.axis_images)\n",
        "\n",
        "    def generate(self):\n",
        "        self.pbar = tqdm.tqdm(total = len(self.frames), position=0, leave=True, desc=\"Building animation... \")\n",
        "        self.anim = animation.FuncAnimation(self.fig, self.animate, frames=len(self.frames), interval=100, blit=True)\n",
        "\n",
        "        # Note: below is the part which makes it work on Colab\n",
        "        rc('animation', html='jshtml')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7tdIkzmNlcMh"
      },
      "source": [
        "## 0.3. Testing your code and grading\n",
        "\n",
        "For your convenience, we have written a helper function for you to write and run test cases for the functions you have to write. One public test case has been provided for each function, but we will also test your code against hidden test cases. For example, if the function signature specifies inputs of size (N,d), you should ensure that your code works for multiple values of N and d.\n",
        "\n",
        "Your code does not need to handle edge cases or improper inputs, and it does not need to run efficiently. Our autograder does not give partial credit for beautiful, well-documented code if it is wrong. If you change any of the function names, the autograder will fail and your score will be 0."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aKp_HqjILbGB"
      },
      "outputs": [],
      "source": [
        "def check_function(function: Callable, test_inputs: list, expected_outputs: list):\n",
        "    \"\"\"Checks that the output of a function with given inputs is close to the expected output value.\"\"\"\n",
        "    actual_outputs = function(*test_inputs)\n",
        "    if isinstance(actual_outputs, tuple):\n",
        "        actual_outputs = list(actual_outputs)\n",
        "    else:\n",
        "        actual_outputs = [actual_outputs]\n",
        "\n",
        "    for i in range(len(expected_outputs)):\n",
        "        assert torch.allclose(actual_outputs[i], expected_outputs[i], rtol=1e-03, atol=1e-05), f\"Output mismatch\\n\\n  Expected:\\n{expected_outputs[i]}\\n\\n  Actual Output:\\n{actual_outputs[i]}\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Vq1rwQbY0dCV"
      },
      "source": [
        "# 1. Coordinate Transforms & Projective Geometry\n",
        "We'll start with writing some simple transformations on points in 3D space, learn how to project them to 2D image space and apply them to render a simple point cloud."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zcXctUr8MLWA"
      },
      "source": [
        "## 1.1. Rigid transformations\n",
        "\n",
        "A rigid-body transform of a 3D point consists of a rotation and a translation. Such transforms are necessary whenever converting between camera coordinates and 3D world coordinates.\n",
        "\n",
        "Given a point $\\mathbf{v} \\in \\mathbb{R}^3$ that we want to transform, a rigid-body transform is formally defined as:\n",
        "$$\n",
        "\\mathbf{v}' = \\mathbf{R} \\mathbf{v} + \\mathbf{t},\n",
        "$$\n",
        "with a rotation matrix $\\mathbf{R} \\in SO(3)$, and a translation $\\mathbf{t} \\in \\mathbb{R}^3$.\n",
        "\n",
        "Consider appending a fourth dimension to $\\mathbf{v}$ with value 1. Then instead of breaking up a rigid body transformation transformation into a matrix-vector product and a sum, we may instead write it as a single matrix-vector product:\n",
        "$$\n",
        "\\mathbf{v}' = \\mathbf{T} \\mathbf{v},\n",
        "$$\n",
        "where\n",
        "$$\n",
        "    \\mathbf{T} = \\begin{bmatrix}\n",
        "                    \\mathbf{R} & \\mathbf{t} \\\\\n",
        "                    \\mathbf{0} & 1\n",
        "                 \\end{bmatrix}\n",
        "$$\n",
        "\n",
        "When $\\mathbf{v}$ has this 4D form, it is called **homogeneous coordinates**. Using homogeneous coordinates is convenient as it allows us to compose multiple coordinate transformations by performing matrix products. As well, we can revert transformations by simply inverting the matrix $\\mathbf{T}$, i.e.:\n",
        "$$\n",
        "    \\mathbf{v} = \\mathbf{T}^{-1} \\mathbf{v}.\n",
        "$$\n",
        "\n",
        "Note that appending 1 to $\\mathbf{v}$ has the effect of adding $\\mathbf{t}$ to the result of the rotation. There are times when we also want to express directional vectors in $\\mathbb{R}^3$ (e.g. normal vectors on a surface) using homogeneous coordinates. Because directional vectors are affected by rotation but not by translation, we append a 0 to such vectors.\n",
        "\n",
        "Homogeneous coordinates can also be used to express 2D rotations and translations, in which case a third dimension is appended.\n",
        "\n",
        "Below, implement three functions: one to homogenize points, one to homogenize directional vectors, and one to apply a rigid-body transform to homogenized points or directions."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FhmcnWRhzs5-",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "# @title 0.5 points\n",
        "def homogenize_points(points: torch.Tensor) -> torch.Tensor:\n",
        "    \"\"\"\n",
        "    Args:\n",
        "        points: points of shape (N, d)\n",
        "\n",
        "    Returns:\n",
        "        vectors of shape (N, d+1)\n",
        "    \"\"\"\n",
        "    raise NotImplementedError\n",
        "\n",
        "# Example test case: homogenize_points([[0, 0, 0]]) should map to [[0, 0, 0, 1]]\n",
        "check_function(homogenize_points,\n",
        "               [torch.zeros(1, 3)],\n",
        "               [torch.tensor([[0., 0., 0., 1.]])])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ogf-8DcXBurV",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "# @title 0.5 points\n",
        "def homogenize_dirs(directions: torch.Tensor) -> torch.Tensor:\n",
        "    \"\"\"\n",
        "    Args:\n",
        "        directions: directional vectors of shape (N, d)\n",
        "\n",
        "    Returns:\n",
        "        vectors of shape (N, d+1)\n",
        "    \"\"\"\n",
        "    raise NotImplementedError\n",
        "\n",
        "check_function(homogenize_dirs,\n",
        "               [torch.zeros(1, 3)],\n",
        "               [torch.zeros(1, 4)])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IsWtdZXICfPR",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "# @title 1 point\n",
        "def transform_rigid(xyz_hom: torch.Tensor, T: torch.Tensor) -> torch.Tensor:\n",
        "    \"\"\"Apply a rigid-body transform to a batch of points / directions.\n",
        "\n",
        "    Args:\n",
        "        xyz_hom: homogeneous 3D vectors of shape (N,4)\n",
        "        T: rigid-body transform matrix of shape (4,4)\n",
        "\n",
        "    Returns:\n",
        "        transformed vectors in homogeneous coordinates\n",
        "    \"\"\"\n",
        "    raise NotImplementedError\n",
        "\n",
        "check_function(transform_rigid,\n",
        "               [torch.tensor([[0., 0, 0, 1]]),\n",
        "                torch.tensor([[0, 0, 0, 1.],\n",
        "                              [0, 0, 0, 0],\n",
        "                              [0, 0, 0, 0],\n",
        "                              [0, 0, 0, 1]])],\n",
        "               [torch.tensor([[1., 0, 0, 1.]])])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Mggi7T2jUGyF"
      },
      "source": [
        "## 1.2. Camera parameters and transforms\n",
        "\n",
        "To relate the image captured by a camera to 3D points in a scene, we model a pinhole camera via two parameters: the extrinsic camera matrix (aka camera pose) and the intrinsic camera matrix.\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "krdZ2_945Mk1"
      },
      "source": [
        "\n",
        "### 1.2. (a) Camera Extrinsics\n",
        "\n",
        "The extrinsic camera matrix (pose) is the rigid transformation that defines the position and rotation of a camera relative to some canonical frame of reference that is referred to as **world coordinates**. By convention, we store the ```cam2world``` format, i.e., a matrix that acts on camera coordinates to produce world coordinates. Naturally, the ```world2cam``` transform is its inverse.\n",
        "\n",
        "For the rest of the course, we adopt the **OpenCV conventions** for camera coordinates. That means:\n",
        "\n",
        "*   Image origin $(0, 0)$ is at top left corner\n",
        "*   y-axis points down\n",
        "*   x-axis points right\n",
        "*   z-axis points into the image plane, i.e., things further in front of the camera have larger z values.\n",
        "\n",
        "Next, we will write convenience functions to transform points and vectors from camera to world coordinates and vice-versa.\n",
        "\n",
        "Below, **implement the functions `transform_world2cam` and `transform_cam2world`**, to be called with homogenized 3D points/vectors, following the convention that all functions will always assume that they will be called with ```cam2world``` matrices.  Make use of the fact that the rigid-body transforms for world2cam and cam2world are inverses of each other, and use the rigid-body transform function you have written above.\n",
        "\n",
        "To further understand camera matrices, you may find it helpful to read this [3 part blog](https://ksimek.github.io/2012/08/14/decompose/)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cxHIYTT1SJUr",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "# @title 1 point\n",
        "def transform_world2cam(xyz_world_hom: torch.Tensor, cam2world: torch.Tensor) -> torch.Tensor:\n",
        "    \"\"\"Transforms points from world coordinates to 3D camera coordinates.\n",
        "\n",
        "    Args:\n",
        "        xyz_world_hom: homogenized 3D points of shape (N, 4)\n",
        "        cam2world: camera pose of shape (4, 4)\n",
        "\n",
        "    Returns:\n",
        "        points in homogeneous camera coordinates (N, 4)\n",
        "    \"\"\"\n",
        "    raise NotImplementedError\n",
        "\n",
        "check_function(transform_world2cam,\n",
        "               [torch.tensor([[0., 0., 0., 1.]]),\n",
        "                torch.tensor([[1., 0., 0., 0.],\n",
        "                            [0., 1., 0., 5.],\n",
        "                            [0., 0., 1., 0.],\n",
        "                            [0., 0., 0., 1.]])],\n",
        "               [torch.tensor([[0., -5., 0., 1]])])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MQ9T8ji7hiKs",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "# @title 1 point\n",
        "def transform_cam2world(xyz_cam_hom: torch.Tensor, cam2world: torch.Tensor) -> torch.Tensor:\n",
        "    \"\"\"Transforms points from 3D camera coordinates to world coordinates.\n",
        "\n",
        "    Args:\n",
        "        xyz_cam_hom: homogenized 3D points of shape (N, 4)\n",
        "        cam2world: camera pose of shape (4, 4)\n",
        "\n",
        "    Returns:\n",
        "        points in homogeneous world coordinates of shape (N, 4)\n",
        "    \"\"\"\n",
        "    raise NotImplementedError\n",
        "\n",
        "check_function(transform_cam2world,\n",
        "               [torch.tensor([[0., 0., 0., 1.]]),\n",
        "                torch.tensor([[1., 0., 0., 0.],\n",
        "                            [0., 1., 0., 5.],\n",
        "                            [0., 0., 1., 0.],\n",
        "                            [0., 0., 0., 1.]])],\n",
        "               [torch.tensor([[0., 5., 0., 1]])])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "h40X3e6Muvvu"
      },
      "source": [
        "### 1.2. (b) Camera Intrinsics\n",
        "\n",
        "Next, we will look at intrinsic camera parameters. As discussed in the lecture, these store the parameters of the projection information. I.e., this is a mapping that takes a point from 3D camera coordinates to the coordinates of that point on the image plane of the camera, i.e., pixel coordinates.\n",
        "\n",
        "For a pinhole camera, the intrinsic camera matrix is of the form\n",
        "$$\n",
        "\\mathbf{K}_{cv} = \\begin{bmatrix}\n",
        "                f_x & 0 & c_x \\\\\n",
        "                0 & f_y & c_y \\\\\n",
        "                0 & 0 & 1 \\\\\n",
        "             \\end{bmatrix}\n",
        "$$\n",
        "\n",
        "with $f_x$ and $f_y$ the focal length for width and height (almost always equal), and $c_x$ and $c_y$ the $x$, $y$ pixel coordinates of the center pixel.\n",
        "\n",
        "Side note: you may recall that the focal length of a lens determines the distance at which a sensor must be placed in order for an object to be in focus. In a pinhole camera all objects are always in focus, so we can place the image plane anywhere focal length merely describes the size of a pixel relative to camera coordinates.\n",
        "\n",
        "Following the lecture, implement the function \"project\" that takes a homogeneous 3D point in camera coordinates as well as camera intrinsics and projects them onto the image plane to return homogeneous points in pixel coordinates."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ulxboDiD4_fG",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "# @title 2 points\n",
        "def project(xyz_cam_hom: torch.Tensor, intrinsics: torch.Tensor) -> torch.Tensor:\n",
        "    \"\"\"Projects camera coordinates to pixel coordinates.\n",
        "\n",
        "    Args:\n",
        "        xyz_cam_hom: 3D points of shape (N, 4)\n",
        "        intrinsics: camera intrinsics of shape (3, 3)\n",
        "\n",
        "    Returns:\n",
        "        homogeneous pixel coordinates of shape (N, 3)\n",
        "    \"\"\"\n",
        "    raise NotImplementedError\n",
        "\n",
        "check_function(project,\n",
        "               [torch.tensor([[0., 0., 5, 1]]),\n",
        "                torch.tensor([[32., 0, 64],\n",
        "                              [0,  32, 64],\n",
        "                              [0,   0,  1]])],\n",
        "               [torch.tensor([[64., 64., 1.]])])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eblV24_zI9PO"
      },
      "source": [
        "## 1.3. Transforming and Rendering Mesh Vertices"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w7RWEnooMOXJ"
      },
      "source": [
        "### Downloading the Bunny Mesh\n",
        "\n",
        "We will use a 3D scan from 1993 of the [famous \"Stanford bunny\"](https://faculty.cc.gatech.edu/~turk/bunny/bunny.html) by Greg Turk as a testbed for the code you'll write. For this first part, we'll only use the vertices of the mesh, which make up a point cloud. We download the mesh from [alecjacobson](http://www.cs.toronto.edu/~jacobson/)'s [`common-3d-test-models`](https://github.com/alecjacobson/common-3d-test-models) repository."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "526V3i90MOAq"
      },
      "outputs": [],
      "source": [
        "url = \"https://raw.githubusercontent.com/alecjacobson/common-3d-test-models/master/data/stanford-bunny.obj\"\n",
        "!wget --quiet --show-progress --no-clobber {url}\n",
        "!ls -l stanford-bunny.obj"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QLPEd8BYzGCP"
      },
      "source": [
        "We use `pytorch3d` to load the mesh, which allows us to get its vertices by calling `.get_mesh_verts_faces()`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "B-HmU8QYahrc"
      },
      "outputs": [],
      "source": [
        "import openmesh\n",
        "bunny_mesh = openmesh.read_trimesh('stanford-bunny.obj')\n",
        "verts = torch.tensor(bunny_mesh.points(), dtype=torch.float)\n",
        "verts = verts - verts.mean(dim=0)\n",
        "verts = verts / verts.max()\n",
        "verts = verts * 1.4"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PuunrBonMpB9"
      },
      "source": [
        "Here is a matplotlib helper function to plot a scatterplot of 2D points in image space."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xL-kFKAWMonp"
      },
      "outputs": [],
      "source": [
        "def plot_image_space_points(\n",
        "    vertices,\n",
        "    alpha=.5,\n",
        "    title=None,\n",
        "    max_points=10_000,\n",
        "    xlim=(-1, 1),\n",
        "    ylim=(-1, 1)):\n",
        "    \"\"\"Plot a pointcloud tensor of shape (N, coordinates)\n",
        "    \"\"\"\n",
        "    vertices = vertices.cpu()\n",
        "\n",
        "    assert len(vertices.shape) == 2\n",
        "    N, dim = vertices.shape\n",
        "    assert dim==2\n",
        "\n",
        "    if N > max_points:\n",
        "        vertices = np.random.default_rng().choice(vertices, max_points, replace=False)\n",
        "\n",
        "    fig = plt.figure(figsize=(6,6))\n",
        "\n",
        "    ax = fig.add_subplot(111)\n",
        "\n",
        "    ax.set_xlabel(\"x\")\n",
        "    ax.set_ylabel(\"y\")\n",
        "\n",
        "    ax.set_xlim(xlim)\n",
        "    ax.set_ylim(ylim)\n",
        "    # ax.set_axis_off()\n",
        "\n",
        "    ax.scatter(*vertices.T, alpha=alpha, marker=',', lw=.5, s=1, color='black')\n",
        "    plt.show(fig)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EXAfgq7AaxaQ"
      },
      "source": [
        "### 1.3. (a) Simple point cloud renderer\n",
        "\n",
        "Time to put everything together to build a simple point cloud renderer!\n",
        "\n",
        "First, let's define a test camera - we'll need to define the camera pose, `cam2world`, and the intrinsics parameters, `K`.\n",
        "\n",
        "Below, we specified for you a `cam2world` matrix and an intrinsics matrix `K` such that:\n",
        "\n",
        "#### Camera pose specifications:\n",
        "* The camera center is at $O = (x_o, y_o, z_o) = (0, 0, 3)$.\n",
        "* The camera rotation is such that if the bunny was rendered, it would be rendered \"Right side up\", i.e., in a frontal view. Hint: The bunny is \"standing\" on a plane parralel to the XZ plane, and the bunny's \"Up\" is in the positive Y-axis. With this information and the knowledge of the OpenCV specification from 1.2. (a), you can figure out the rotation matrix part of the cam2world matrix.\n",
        "\n",
        "#### Camera intrinsics specifications:\n",
        "* The image should have a sidelength of 1, with the pixel center at 0.5.\n",
        "* The field of view should be 90 degree.\n",
        "* You can then figure out the focal length via the equation $\\text{tan}(0.5*\\text{FOV}) = (0.5 * \\text{image_size}) / (\\text{focal_length})$"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QbbW8c0TOZzu"
      },
      "outputs": [],
      "source": [
        "img_resolution = (256, 256)\n",
        "\n",
        "test_cam2world = torch.tensor([[1., 0., 0., 0.],\n",
        "                               [0., -1., 0., 0.],\n",
        "                               [0., 0., -1., 3.],\n",
        "                               [0., 0., 0., 1.]], device=device)\n",
        "test_K  = torch.tensor([[0.5, 0., 0.5],\n",
        "                        [0., 0.5, 0.5],\n",
        "                        [0., 0.,  1.]], device=device)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Grj0KARCAmqB"
      },
      "source": [
        "We will now project the points onto the camera image plane and do a scatterplot of that!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YLbKqP7Po3mW"
      },
      "outputs": [],
      "source": [
        "v_hom = homogenize_points(verts)\n",
        "xyz_cam_hom = transform_world2cam(v_hom, test_cam2world)\n",
        "xy_pix_hom = project(xyz_cam_hom, test_K)\n",
        "\n",
        "plot_image_space_points(xy_pix_hom[..., :2], xlim=(0, 1), ylim=(0, 1))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EKzA97YKBNuB"
      },
      "source": [
        "We note a few things:\n",
        "\n",
        "1.   The bunny is on its head.\n",
        "2.   This is a scatterplot of the continuous coordinates of projected points, whereas an image is usually discretized into a grid of pixels.\n",
        "3.   Note that the origin $(0, 0)$ of this plot is at the bottom left! This is different from the OpenCV camera coordinate system that we assumed, where the origin is at the top left!\n",
        "\n",
        "We have provided a function, `discretize_normalized_2d_pointclouid_to_img`, that addresses these points. You do not need to modify it, simply use it to display your point cloud scatterplot."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "S30gKBggOtjk"
      },
      "outputs": [],
      "source": [
        "def discretize_normalized_2d_pointcloud_to_img(\n",
        "    pointcloud_2d,\n",
        "    img_resolution = (100, 100)\n",
        "):\n",
        "    img = np.zeros(img_resolution)\n",
        "    x, y = np.split(pointcloud_2d, 2, axis=-1)\n",
        "    x *= img_resolution[0]\n",
        "    y *= img_resolution[1]\n",
        "\n",
        "    x, y = x.int(), y.int()\n",
        "    img[y, x] = 1.\n",
        "\n",
        "    return img\n",
        "\n",
        "img = discretize_normalized_2d_pointcloud_to_img(xy_pix_hom.cpu()[..., :2],\n",
        "                                          img_resolution=img_resolution)\n",
        "\n",
        "plt.figure(figsize=(10, 10))\n",
        "plt.imshow(img, cmap='Greys')\n",
        "plt.axis(\"off\");\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ot149_Owxnox"
      },
      "source": [
        "You should see an upright bunny silhouette. There may be a few empty pixels in the middle. This is normal since the vertices are not particularly dense.\n",
        "\n",
        "To finish up Problem 1, let's put `transform_rigid()` to use and rotate the bunny.\n",
        "\n",
        "Fill in `rotate_by_theta()` which rotates input 3D points about the **y-axis** in the clockwise direction (when viewed from above). You will have to come up with the appropriate rotation matrix and apply it using `transform_rigid()`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ek2XGPvpO5-B"
      },
      "outputs": [],
      "source": [
        "# @title 1 point\n",
        "def rotate_by_theta(xyz: torch.Tensor, theta: torch.Tensor) -> torch.Tensor:\n",
        "    \"\"\"Rotate a (batch of) points around the origin.\n",
        "\n",
        "    Args:\n",
        "        xyz: 3D points of shape (N, 3)\n",
        "        theta: scalar angle value in degrees. Shape: ()\n",
        "\n",
        "    Returns:\n",
        "        transformed points of shape (N, 3)\n",
        "    \"\"\"\n",
        "\n",
        "    raise NotImplementedError\n",
        "\n",
        "check_function(rotate_by_theta,\n",
        "               [torch.tensor([[0., 0., 1]]), torch.tensor(90.)],\n",
        "               [torch.tensor([[-1., 0., 0.]])])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "clRgq73mVfNc"
      },
      "source": [
        "We'll use the rotation method to animate a rotating bunny!\n",
        "You don't have to modify `animate_bunny()`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hnwQwp8zTdad"
      },
      "outputs": [],
      "source": [
        "def animate_bunny(verts, img_resolution=(100, 100)):\n",
        "    frames = []\n",
        "    for theta in range(0, 360, 10):\n",
        "        rotated_verts = rotate_by_theta(verts, torch.tensor(theta))\n",
        "        v_hom = homogenize_points(rotated_verts)\n",
        "        xyz_cam_hom = transform_world2cam(v_hom, test_cam2world)\n",
        "        xy_pix_hom = project(xyz_cam_hom, test_K)\n",
        "\n",
        "        image = discretize_normalized_2d_pointcloud_to_img(\n",
        "                    xy_pix_hom.cpu()[..., :2],\n",
        "                    img_resolution=img_resolution)\n",
        "        frames.append([image])\n",
        "\n",
        "    return AnimationWrapper(rows=1, cols=1, frames=frames)\n",
        "\n",
        "rotating_bunny_animation = animate_bunny(verts)\n",
        "rotating_bunny_animation.generate()\n",
        "rotating_bunny_animation.anim"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "l9kfcofXV9Am"
      },
      "source": [
        "You should see a bunny silhouette rotating clockwise about the y-axis ('Up' direction)\n",
        "\n",
        "If you do see this, your implementation of the transformation methods are likely correct! Proceed to the next problem."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fniw72sYzJBR"
      },
      "source": [
        "# 2. Structured Light\n",
        "\n",
        "We'll now apply these transformations to the problem of recovering depth using structured light. A real-world example of a structured light setup is the Kinect.\n",
        "\n",
        "The Kinect consists of a camera and a projector, as shown in the figure below. The projector shines a structured pattern of light on the world (in the infrared spectrum so that it is invisible to our eyes).\n",
        "The camera can see this pattern and observes how the light warps over the scene geometry.\n",
        "\n",
        "The geometry of light rays emitted by the projector can be modeled similarly to a pinhole camera -- the projector has its own pose, intrinsics, image plane, etc."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Q2zEdA14WBa6"
      },
      "source": [
        "![](https://drive.google.com/uc?id=1fA-L1Sf5QnCj9J0VJFcc4zbBxmjrL8v4)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fi8MyccbyL_u"
      },
      "source": [
        "## 2.1. Simulating a Light Projector\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wFDkiE4lwY8a"
      },
      "source": [
        "### Downloading the Motorcycle Scene\n",
        "\n",
        "We'll download and process the motorcycle scene here. You do not have to modify this cell."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cC1VskHhMbRO"
      },
      "outputs": [],
      "source": [
        "# Before starting, run this cell to download the motorbike scene that we'll be using.\n",
        "import scipy.io as sio\n",
        "import requests\n",
        "import cv2\n",
        "import os\n",
        "\n",
        "def download(url, fn):\n",
        "    r = requests.get(url)\n",
        "    if r.status_code == 200:\n",
        "        open(fn, 'wb').write(r.content)\n",
        "        print(\"{} downloaded: {:.2f} KB\".format(fn, len(r.content)/1024.))\n",
        "    else:\n",
        "        print(\"url not found:\", url)\n",
        "\n",
        "os.makedirs('Motorcycle-perfect', exist_ok=True)\n",
        "download('http://6.869.csail.mit.edu/sp22/pset2_data/Motorcycle-perfect/rgb_camera_view.png', 'Motorcycle-perfect/rgb_camera_view.png')\n",
        "download('http://6.869.csail.mit.edu/sp22/pset2_data/Motorcycle-perfect/stripe_lights.png', 'Motorcycle-perfect/stripe_lights.png')\n",
        "download('http://6.869.csail.mit.edu/sp22/pset2_data/Motorcycle-perfect/z_camera_view.mat', 'Motorcycle-perfect/z_camera_view.mat')\n",
        "download('http://6.869.csail.mit.edu/sp22/pset2_data/Motorcycle-perfect/z_projector_view.mat', 'Motorcycle-perfect/z_projector_view.mat')\n",
        "\n",
        "def make_motorcycle_scene(scale=0.2):\n",
        "\n",
        "    f0 = 3979.911  # focal length in (pixels)\n",
        "    d = 193.001  # baseline in (mm)\n",
        "    cx_c, cy_c = (1244.772, 1019.507)  # principal point coordinates for the camera\n",
        "    cx_p, cy_p = (1369.115, 1019.507)\n",
        "\n",
        "    Z_p_img = sio.loadmat(\"./Motorcycle-perfect/z_projector_view.mat\")[\"z_p_img\"]\n",
        "    Z_c_img = sio.loadmat(\"./Motorcycle-perfect/z_camera_view.mat\")[\"z_c_img\"]\n",
        "    rgb_c_img = cv2.imread(\"./Motorcycle-perfect/rgb_camera_view.png\")\n",
        "\n",
        "    ksize = int(1.0 / scale) * 2 + 1\n",
        "    Z_p_img = cv2.GaussianBlur(Z_p_img, (ksize, ksize), 0)\n",
        "    Z_p_img = cv2.resize(Z_p_img, (0, 0), fx=scale, fy=scale)\n",
        "\n",
        "    Z_c_img = cv2.GaussianBlur(Z_c_img, (ksize, ksize), 0)\n",
        "    Z_c_img = cv2.resize(Z_c_img, (0, 0), fx=scale, fy=scale)\n",
        "\n",
        "    rgb_c_img = cv2.GaussianBlur(rgb_c_img, (ksize, ksize), 0)\n",
        "    rgb_c_img = cv2.resize(rgb_c_img, (0, 0), fx=scale, fy=scale)\n",
        "\n",
        "    f = f0 * scale\n",
        "    cx_c = cx_c * scale\n",
        "    cy_c = cy_c * scale\n",
        "    cx_p = cx_p * scale\n",
        "    cy_p = cy_p * scale\n",
        "\n",
        "    img_size = Z_p_img.shape\n",
        "\n",
        "    proj_intrinsics = torch.tensor([[0.,  f,  cy_p],\n",
        "                                    [f,  0.,  cx_p],\n",
        "                                    [0., 0.,   1.]], device=device)\n",
        "\n",
        "    cam_intrinsics = torch.tensor([[0.,  f,  cy_c],\n",
        "                                  [f,  0.,  cx_c],\n",
        "                                  [0., 0.,   1.]], device=device)\n",
        "\n",
        "    proj_extrinsics = torch.tensor([[1.0,  0.,  0., 0.],\n",
        "                                    [0.,  -1.,  0., 0.],\n",
        "                                    [0.,   0., -1., 0.],\n",
        "                                    [0.,   0.,  0., 1.]], device=device)\n",
        "\n",
        "    cam_extrinsics = torch.tensor([[1.0,  0.,  0., -d],\n",
        "                                  [0.,  -1.,  0., 0.],\n",
        "                                  [0.,   0., -1., 0.],\n",
        "                                  [0.,   0.,  0., 1.]], device=device)\n",
        "\n",
        "    return Scene(rgb_ground_truth=torch.tensor(rgb_c_img),\n",
        "                z_proj_ground_truth=torch.tensor(Z_p_img).float(),\n",
        "                z_cam_ground_truth=torch.tensor(Z_c_img).float(),\n",
        "                cam_extrinsics=cam_extrinsics.float(),\n",
        "                cam_intrinsics=cam_intrinsics.float(),\n",
        "                proj_extrinsics=proj_extrinsics.float(),\n",
        "                proj_intrinsics=proj_intrinsics.float())\n",
        "\n",
        "motorcycle_scene = make_motorcycle_scene(scale=0.04)\n",
        "motorcycle_scene_large = make_motorcycle_scene(scale=0.2)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "H2dZ78HxtZmB"
      },
      "source": [
        "We first try to simulate a laser projector which supplies illumination for a single projector pixel. i.e. for a given point on the projector's image space `uv_p`, find the corresponding point on the camera's image space that will be lit up.\n",
        "\n",
        "We'll do this in these steps:\n",
        "\n",
        "\n",
        "1.   'unproject' `uv_p` to get projector-space `xy_p` (3D point relative to projector's origin)\n",
        "2.   Apply `cam2world()` using the _projector_'s extrinsic matrix to get `xy_w` (3D point relative to world's origin)\n",
        "3.   Apply `world2cam()` using the _camera_'s extrinsic matrix to get `xy_c` (3D point relative to camera's origin)\n",
        "4.   Use `project()` with _camera_'s intrinsic matrix to get `uv_c` (2D point on the camera's image-space)\n",
        "\n",
        "We've already constructed `cam2world()`, `world2cam()` and `project()` in Problem 1.\n",
        "\n",
        "The missing piece here is `unproject()`. You next task is to fill that in:\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ErKGnfgCtZ2v"
      },
      "source": [
        "### 2.1. (a) The \"unproject\" Function\n",
        "\n",
        "You'll now write a function ```unproject``` that takes x, y pixel coordinates and lifts them to 3D camera coordinates.\n",
        "\n",
        "Recall the following equation for projecting a 3D coordinate in camera coordinates $\\mathbf{X}$ to _homogeneous_ pixel coordinates:\n",
        "$$\n",
        "\\begin{align}\n",
        "    \\begin{pmatrix}\n",
        "    f \\cdot X + Z \\cdot p_x \\\\\n",
        "    f \\cdot Y + Z \\cdot p_y\\\\\n",
        "    Z\n",
        "    \\end{pmatrix} &= \\mathbf{K} [\\mathbf{I} | \\mathbf{0} ]\n",
        "    \\begin{pmatrix}\n",
        "    X \\\\\n",
        "    Y \\\\\n",
        "    Z \\\\\n",
        "    1\n",
        "    \\end{pmatrix} \\\\\n",
        "    \\begin{pmatrix}\n",
        "    f \\cdot X + Z \\cdot p_x \\\\\n",
        "    f \\cdot Y + Z \\cdot p_y\\\\\n",
        "    Z\n",
        "    \\end{pmatrix} &= \\mathbf{K}\n",
        "    \\begin{pmatrix}\n",
        "    X \\\\\n",
        "    Y \\\\\n",
        "    Z\n",
        "    \\end{pmatrix}\n",
        "\\end{align}\n",
        "$$\n",
        "\n",
        "Applying the inverse $\\mathbf{K}^{-1}$ on both sides, we see:\n",
        "\n",
        "$$\n",
        "\\begin{align}\n",
        "    \\begin{pmatrix}\n",
        "    X \\\\\n",
        "    Y \\\\\n",
        "    Z\n",
        "    \\end{pmatrix} &= \\mathbf{K}^{-1}\n",
        "    \\begin{pmatrix}\n",
        "    f \\cdot X + Z \\cdot p_x \\\\\n",
        "    f \\cdot Y + Z \\cdot p_y\\\\\n",
        "    Z\n",
        "    \\end{pmatrix} \\\\\n",
        "    \\begin{pmatrix}\n",
        "    X \\\\\n",
        "    Y \\\\\n",
        "    Z\n",
        "    \\end{pmatrix} &= Z \\cdot \\mathbf{K}^{-1}\n",
        "    \\begin{pmatrix}\n",
        "    \\frac{f \\cdot X}{Z} + p_x \\\\\n",
        "    \\frac{f \\cdot Y}{Z} + p_y\\\\\n",
        "    1\n",
        "    \\end{pmatrix} \\\\\n",
        "    \\mathbf{X} &= Z \\cdot \\mathbf{K}^{-1}\n",
        "    \\begin{pmatrix}\n",
        "    x_{pix} \\\\\n",
        "    y_{pix} \\\\\n",
        "    1\n",
        "    \\end{pmatrix}\n",
        "\\end{align}\n",
        "$$\n",
        "\n",
        "This relates 3D world coordinates along a ray to their pixel coordinates and the depth, and is exactly what you have to implement below."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3wnOlpMltr33"
      },
      "outputs": [],
      "source": [
        "# @title 2 points\n",
        "def unproject(\n",
        "    uv_pix: torch.Tensor,\n",
        "    z: torch.Tensor,\n",
        "    intrinsics: torch.Tensor\n",
        "    ) -> torch.Tensor:\n",
        "    \"\"\"Unproject (lift) pixel coordinates x_pix and per-pixel z coordinate\n",
        "    to camera coordinates.\n",
        "\n",
        "    Args:\n",
        "        uv_pix: pixel coordinates of shape (N, 2)\n",
        "        z: per-pixel depth, defined as z coordinate of shape (N,)\n",
        "        intrinsics: camera intrinsics of shape (3, 3)\n",
        "\n",
        "    Returns:\n",
        "        points in 3D camera coordinates.\n",
        "    \"\"\"\n",
        "\n",
        "    raise NotImplementedError\n",
        "\n",
        "check_function(unproject,\n",
        "                [torch.tensor([[0.0, 0.0]], device=device),\n",
        "                 torch.tensor([5.0], device=device),\n",
        "                 torch.tensor([[0.5, 0.0, 0.5],\n",
        "                               [0.0, 0.5, 0.5],\n",
        "                               [0.0, 0.0, 1.0]], device=device)],\n",
        "                [torch.tensor([[-5.0000, -5.0000, 5.0000]]).to(device)])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZAONfMnTwpWw"
      },
      "source": [
        "### 2.1. (b) Simulating Structured Light\n",
        "Now that we have all 4 pieces, let's write a function, `map_projector_coord_to_camera_coord`, that will take an projector's image space `uv_p` and map it to the camera's image space coordinate `uv_c`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4H7GQxhaxDM1"
      },
      "outputs": [],
      "source": [
        "# @title 3 points\n",
        "def map_projector_coord_to_camera_coord(\n",
        "    uv_p: torch.Tensor,\n",
        "    z_p: torch.Tensor,\n",
        "    proj_extrinsics: torch.Tensor,\n",
        "    proj_intrinsics: torch.Tensor,\n",
        "    cam_extrinsics: torch.Tensor,\n",
        "    cam_intrinsics: torch.Tensor\n",
        "  ) -> torch.Tensor:\n",
        "    \"\"\"Map the coordinate on the projector's image space to the\n",
        "      camera's image space by following the 4 steps laid out in 2.1\n",
        "\n",
        "     Args:\n",
        "        uv_p: 2D pixel coordinates of shape (N, 2)\n",
        "        z_p: corresponding pixel depths, defined as z coordinate of shape (N,)\n",
        "        proj_extrinsics: projector's extrinsics of shape (4, 4)\n",
        "        proj_intrinsics: projector's intrinsics of shape (3, 3)\n",
        "        cam_extrinsics: camera's extrinsics of shape (4, 4)\n",
        "        cam_intrinsics: camera's intrinsics of shape (3, 3)\n",
        "\n",
        "     Returns: tuple(uv_c, z_c) where\n",
        "        uv_c: 2D pixel coordinates of shape (N, 2)\n",
        "        z_c: corresponding depths as seen by the camera (N,)\n",
        "    \"\"\"\n",
        "\n",
        "    #\n",
        "    # Try to follow these steps to complete this task:\n",
        "    #\n",
        "    # 1. Unproject projector's image space uv_p to get xy_p\n",
        "    #    (3D coordinates relative to projector origin)\n",
        "    #\n",
        "    # 2. Transform xy_p to world coordinates (xy_w).\n",
        "    #\n",
        "    # 3. Transform xy_w to camera coordinates (xy_c).\n",
        "    #    (3D cooridnates relative to camera origin)\n",
        "    #\n",
        "    # 4. Project onto camera's image space.\n",
        "    #\n",
        "    # (You may need to homogenize/dehomogenize your points before\n",
        "    #  applying certain transforms)\n",
        "    #\n",
        "\n",
        "    raise NotImplementedError\n",
        "\n",
        "check_function(map_projector_coord_to_camera_coord,\n",
        "              [torch.tensor([[1.0, 1.0]], device=device),\n",
        "\n",
        "               torch.tensor([2.0], device=device),\n",
        "\n",
        "               torch.tensor([[1., 0., 0., 0. ],\n",
        "                             [0., -1., 0., 0.],\n",
        "                             [0., 0., -1., 4.],\n",
        "                             [0., 0., 0., 1. ]], device=device),\n",
        "\n",
        "               torch.tensor([[0.5, 0.0, 0.5],\n",
        "                             [0.0, 0.5, 0.5],\n",
        "                             [0.0, 0.0, 1.0]], device=device),\n",
        "\n",
        "               torch.tensor([[1., 0., 0.,  2.],\n",
        "                             [0., -1., 0., 0.],\n",
        "                             [0., 0., -1., 4.],\n",
        "                             [0., 0., 0., 1.]], device=device),\n",
        "\n",
        "               torch.tensor([[0.5, 0.0, 0.5],\n",
        "                             [0.0, 0.5, 0.5],\n",
        "                             [0.0, 0.0, 1.0]], device=device)],\n",
        "\n",
        "              [torch.tensor([[0.5, 1.0]], device=device),\n",
        "\n",
        "               torch.tensor([2.0], device=device)])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TNoYRQyLfFqD"
      },
      "source": [
        "Next, we have provided a utility method `render_scene_with_projected_light()` which will use your `map_projector_coord_to_camera_coord()` method on each projector pixel coordinate `uv_p` to find the `uv_c` on the camera's image, and then accumulate the light values in that location.\n",
        "\n",
        "Note that there can be multiple projector coordinates that map to the same camera coordinates (in this case, our utility method takes the depth into account to figure out which light value should used).\n",
        "\n",
        "You do not need to modify this method."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vq6tQwN44i4W"
      },
      "outputs": [],
      "source": [
        "def render_scene_with_projected_light(\n",
        "    light_image,\n",
        "    z_proj,\n",
        "    z_cam_gt,\n",
        "    proj_extrinsics,\n",
        "    proj_intrinsics,\n",
        "    cam_extrinsics,\n",
        "    cam_intrinsics):\n",
        "    \"\"\"\n",
        "      light_image: per-pixel light intensity from the projector. Shape: (W, H)\n",
        "      z_proj: per-pixel depths as seen from the projector. Shape: (W, H)\n",
        "      proj_extrinsics: projector's extrinsics of shape (4, 4)\n",
        "      proj_intrinsics: projector's intrinsics of shape (3, 3)\n",
        "      cam_extrinsics: camera's extrinsics of shape (4, 4)\n",
        "      cam_intrinsics: camera's intrinsics of shape (3, 3)\n",
        "\n",
        "      Returns:\n",
        "        image of the scene rendered with the provided structured light_image. Shape: (W, H)\n",
        "    \"\"\"\n",
        "    assert (isinstance(light_image, torch.Tensor))\n",
        "\n",
        "    W, H = light_image.shape\n",
        "\n",
        "    # Make a set of inputs.\n",
        "    (u_p, v_p) = torch.meshgrid(\n",
        "                        torch.arange(0, W),\n",
        "                        torch.arange(0, H),\n",
        "                        indexing='ij')\n",
        "\n",
        "    uv_p = torch.stack([u_p, v_p], axis=-1).reshape((-1, 2))\n",
        "    z_proj = z_proj.flatten()\n",
        "\n",
        "    light_vals = light_image.flatten()\n",
        "\n",
        "    # To save on compute, only go through non-zero pixels in\n",
        "    # the light image.\n",
        "    #\n",
        "    non_zero_indices = light_vals.nonzero()[:, 0]\n",
        "\n",
        "    light_vals = light_vals[non_zero_indices]\n",
        "    uv_p = uv_p[non_zero_indices, :]\n",
        "    z_proj = z_proj[non_zero_indices]\n",
        "\n",
        "    uv_c, z_cs = map_projector_coord_to_camera_coord(\n",
        "        uv_p,\n",
        "        z_proj,\n",
        "        proj_extrinsics,\n",
        "        proj_intrinsics,\n",
        "        cam_extrinsics,\n",
        "        cam_intrinsics)\n",
        "\n",
        "    rendered_image = torch.zeros((W, H))\n",
        "    num_samples = torch.zeros((W, H))\n",
        "\n",
        "    log_progress = (light_vals.shape[0] > 500)\n",
        "    pbar = tqdm.tqdm(\n",
        "            total=light_vals.shape[0],\n",
        "            position=0,\n",
        "            leave=True,\n",
        "            desc=\"Rendering projected light\") if log_progress else None\n",
        "\n",
        "    uv_c = uv_c.round().int()\n",
        "\n",
        "    # Set depth error to 1% of mean absolute GT depth.\n",
        "    abs_depth_buffer_error = 0.01 * (z_cam_gt).mean()\n",
        "\n",
        "    for ((u_c, v_c), z_c, light_val) in zip(uv_c, z_cs, light_vals):\n",
        "        pbar.update(1) if log_progress else None\n",
        "        if W > u_c >= 0 and H > v_c >= 0 and abs(z_cam_gt[u_c, v_c] - z_c) < abs_depth_buffer_error:\n",
        "          rendered_image[u_c, v_c] += light_val\n",
        "          num_samples[u_c, v_c] += 1\n",
        "\n",
        "    pbar.close() if log_progress else None\n",
        "\n",
        "    rendered_image = rendered_image / num_samples\n",
        "    rendered_image[rendered_image.isnan()] = 0.0\n",
        "\n",
        "    return rendered_image\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qS56XR8xhc2p"
      },
      "source": [
        "Now, let's see this in action, by shining a stripes pattern from our projector, and viewing our camera image on the side."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OuiTyhWmy9RM"
      },
      "outputs": [],
      "source": [
        "# Load striped lights (single-channel)\n",
        "stripe_lights_image = cv2.imread(\"./Motorcycle-perfect/stripe_lights.png\")[:, :, 0]\n",
        "\n",
        "W, H, C = motorcycle_scene_large.rgb_ground_truth.shape\n",
        "stripe_lights_image = cv2.resize(stripe_lights_image, (H, W))\n",
        "\n",
        "# Render scene with projected light\n",
        "cam_image = render_scene_with_projected_light(\n",
        "    torch.tensor(stripe_lights_image),\n",
        "    motorcycle_scene_large.z_proj_ground_truth,\n",
        "    motorcycle_scene_large.z_cam_ground_truth,\n",
        "    motorcycle_scene_large.proj_extrinsics,\n",
        "    motorcycle_scene_large.proj_intrinsics,\n",
        "    motorcycle_scene_large.cam_extrinsics,\n",
        "    motorcycle_scene_large.cam_intrinsics)\n",
        "\n",
        "# Plot the images\n",
        "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 5))\n",
        "\n",
        "ax1.imshow(cam_image.numpy(), cmap=\"gray\")\n",
        "ax1.set_title('Camera View')\n",
        "\n",
        "ax2.imshow(stripe_lights_image, cmap=\"gray\")\n",
        "ax2.set_title('Projector Image')\n",
        "\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0LSY6KyrhAoZ"
      },
      "source": [
        "You should see the straight stripes in the projector's image turn wiggly as they are projected onto the scene."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ljU-WJx0hpGy"
      },
      "source": [
        "## 2.2 Depth from Single-Pixel Laser\n",
        "\n",
        "Thus far we've been able to map projector coordinates to camera coordinates, *given the true depth* of the scene.\n",
        "\n",
        "Since we know this mapping depends on the depth, what if we could come up with the *inverse* of this operation?\n",
        "Can we infer the depth given a pair of matching coordinates? Yes! This is referred to as **triangulation**, and we will implement a simplified version of this.\n",
        "\n",
        "We need two parts to accomplish this:\n",
        "1. Given a matching pair of projector and camera coordinates `uv_p` and `uv_c`, infer the depth (`infer_depth_from_matched_coords()`)\n",
        "2. A method to generate light patterns (`light_generator_fn()`) in such a way that we can _decode_ one or more pairs of matching coordinates from the camera's view of this projected pattern (`decoder_fn()`).\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f9GH_d8ac3Db"
      },
      "source": [
        "### 2.2 (a) Inferring Depth from Coordinate Pairs\n",
        "\n",
        "\n",
        "Your task is to complete the method `infer_depth_from_matched_coords()`.\n",
        "\n",
        "We will follow a process similar to the lecture, and take advantage of the properties of similar triangles. Here's a diagram of what this looks like:\n",
        "\n",
        "![](https://drive.google.com/uc?id=1oQKFGak8iiEG0DlEuzKeHDBvJUovvLvy)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7b6Bh3dLtCY4"
      },
      "outputs": [],
      "source": [
        "# @title 3 points\n",
        "def infer_depth_from_matched_coords(\n",
        "    uv_c: torch.Tensor,\n",
        "    uv_p: torch.Tensor,\n",
        "    proj_extrinsics: torch.Tensor,\n",
        "    proj_intrinsics: torch.Tensor,\n",
        "    cam_extrinsics: torch.Tensor,\n",
        "    cam_intrinsics: torch.Tensor\n",
        "    ) -> torch.Tensor:\n",
        "    \"\"\"\n",
        "      Args:\n",
        "        uv_c: camera image positions in pixel coordinates of shape (N, 2)\n",
        "        uv_p: projector image positions in pixel coordinates of shape (N, 2)\n",
        "        proj_extrinsics: projector's extrinsic matrix of shape (4, 4)\n",
        "        proj_intrinsics: projector's intrinsic matrix of shape (3, 3)\n",
        "        cam_extrinsics: camera's extrinsic matrix of shape (4, 4)\n",
        "        cam_intrinsics: camera's intrinsic matrix of shape (3, 3)\n",
        "\n",
        "      Returns:\n",
        "        z_c: inferred depths w.r.t the camera's origin (N,)\n",
        "    \"\"\"\n",
        "\n",
        "    raise NotImplementedError\n",
        "\n",
        "check_function(infer_depth_from_matched_coords,\n",
        "                [torch.tensor([[1.0, 1.0], [1.0, 0.6]], device=device),\n",
        "\n",
        "                torch.tensor([[1.0, 0.5], [1.0, 0.5]], device=device),\n",
        "\n",
        "                torch.tensor([[1., 0., 0.,  0. ],\n",
        "                              [0., 1., 0., 0.],\n",
        "                              [0., 0., 1., 4.],\n",
        "                              [0., 0., 0.,  1. ]], device=device),\n",
        "\n",
        "                torch.tensor([[0.0, 0.5, 0.5],\n",
        "                              [0.5, 0.0, 0.5],\n",
        "                              [0.0, 0.0, 1.0]], device=device),\n",
        "\n",
        "                torch.tensor([[1., 0., 0.,  -2.],\n",
        "                              [0., 1., 0., 0.],\n",
        "                              [0., 0., 1., 4.],\n",
        "                              [0., 0., 0., 1.]], device=device),\n",
        "\n",
        "                torch.tensor([[0.0, 0.5, 0.5],\n",
        "                              [0.5, 0.0, 0.5],\n",
        "                              [0.0, 0.0, 1.0]], device=device)],\n",
        "\n",
        "                [torch.tensor([2.0, 10.0], device=device)])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ke3GnFy_ycV1"
      },
      "source": [
        "### Pattern Generator & Decoder Functions for a Laser\n",
        "The last thing we need is a pair of methods: one that generates structured light patterns, and one that recovers pairs of matching coordinates given the image seen by the camera.\n",
        "\n",
        "Remember that our goal is to generate light patterns in such a way that it is easy to recover corresponding pairs of coordinates.\n",
        "\n",
        "For this example, we're going provide an implementation of the simplest answer to this: every projector image has a single pixel lit up (i.e. like a laser pointer). That way, there's at most one pixel illuminated in the camera's image, and there's no ambiguity.\n",
        "\n",
        "We provide the implementations of the following functions for you. You do not have to modify these methods.\n",
        "\n",
        "1.   `per_pixel_laser_light_fn()`: Generates a series of images for the projector with a single non-zero pixel each. We use python _generators_ for this. More info on generators [here](https://www.learnpython.org/en/Generators)\n",
        "\n",
        "2.   `laser_decoder_fn()`: Generates two vectors of matched coordinates given the light image, and the corresponding simulated camera image. In the case of a single-pixel laser, we can only recover a single set of corresponding set of coordinates for each time we use the laser\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LCYZV3A-0HD2"
      },
      "outputs": [],
      "source": [
        "def per_pixel_laser_light(W, H):\n",
        "    (u_ps, v_ps) = torch.meshgrid(\n",
        "          torch.arange(0, W, step=1),\n",
        "          torch.arange(0, H, step=1),\n",
        "          indexing='xy')\n",
        "\n",
        "    uv_p = torch.stack([u_ps, v_ps], axis=-1).reshape((-1, 2))\n",
        "\n",
        "    for (u_p, v_p) in uv_p:\n",
        "        light_image = torch.zeros((W, H))\n",
        "\n",
        "        light_image[u_p, v_p] = 1\n",
        "\n",
        "        yield (light_image, None)\n",
        "\n",
        "def laser_decoder_fn(light_image, camera_image, extra_info=None):\n",
        "  # Since we only expect a single pixel to be lit up, we'll\n",
        "  # just use 'nonzero' to get the coordinates of that pixel.\n",
        "  #\n",
        "  uv_p = light_image.nonzero()\n",
        "  uv_c = camera_image.nonzero()\n",
        "\n",
        "  if (uv_c.shape[0] == 0):\n",
        "    return None, None\n",
        "\n",
        "  return uv_p, uv_c\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y5LZgfUv2HGe"
      },
      "source": [
        "### Putting these Together\n",
        "\n",
        "Now we can put all these pieces together to recover a depth map!\n",
        "\n",
        "We have provided a method `recover_depth_map()` that takes in a pair of structured light generator and decoder methods. You do not have to modify this.\n",
        "\n",
        "Simply run the two cells below to see your results. If your implementation of\n",
        "`infer_depth_from_matched_coordinates` is correct, you will see a good depth recovery animation."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xQwoylZn0non"
      },
      "outputs": [],
      "source": [
        "def recover_depth_map(\n",
        "    W, H,\n",
        "    structured_light_fn,\n",
        "    decoder_fn,\n",
        "    z_proj_gt,\n",
        "    z_cam_gt,\n",
        "    rgb_gt,\n",
        "    proj_extrinsics,\n",
        "    proj_intrinsics,\n",
        "    cam_extrinsics,\n",
        "    cam_intrinsics,\n",
        "    plot_skip = 1\n",
        "    ):\n",
        "\n",
        "    camera_depth_image = torch.ones((W, H)) * torch.nan\n",
        "\n",
        "    frames = []\n",
        "\n",
        "    pbar = tqdm.tqdm(position=0, leave=True, desc=\"Recovering depth with light image\")\n",
        "    for index, (light_image, extra_info) in enumerate(structured_light_fn(W, H)):\n",
        "        simulated_camera_image = render_scene_with_projected_light(\n",
        "                                    light_image,\n",
        "                                    z_proj_gt,\n",
        "                                    z_cam_gt,\n",
        "                                    proj_extrinsics,\n",
        "                                    proj_intrinsics,\n",
        "                                    cam_extrinsics,\n",
        "                                    cam_intrinsics)\n",
        "\n",
        "        uv_p, uv_c = decoder_fn(light_image, simulated_camera_image, extra_info)\n",
        "\n",
        "        if (uv_c is None):\n",
        "          continue\n",
        "\n",
        "        z_cam = infer_depth_from_matched_coords(\n",
        "                    uv_c,\n",
        "                    uv_p,\n",
        "                    proj_extrinsics,\n",
        "                    proj_intrinsics,\n",
        "                    cam_extrinsics,\n",
        "                    cam_intrinsics)\n",
        "\n",
        "        u_c, v_c = uv_c.permute((1, 0))\n",
        "        camera_depth_image[u_c, v_c] = z_cam\n",
        "\n",
        "        if (index % plot_skip == 0):\n",
        "          pbar.update(plot_skip)\n",
        "          frames.append([rgb_gt.numpy(), light_image.numpy(), simulated_camera_image.numpy(), camera_depth_image.clone().numpy()])\n",
        "\n",
        "    depth_animation = AnimationWrapper(rows=2, cols=2, frames=frames)\n",
        "    axes = depth_animation.get_axes()\n",
        "    axes[0].set_title(\"Camera view with natural lighting\")\n",
        "    axes[1].set_title(\"Projector view of projected light\")\n",
        "    axes[2].set_title(\"Camera view of projected light\")\n",
        "    axes[3].set_title(\"Recovered depth map\")\n",
        "\n",
        "    return (camera_depth_image, depth_animation)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9JK0moTR2sp5"
      },
      "outputs": [],
      "source": [
        "W, H = motorcycle_scene.z_proj_ground_truth.shape\n",
        "\n",
        "# Recover a depth map using the laser light generator and decoder method\n",
        "z_recovered, depth_animation = recover_depth_map(\n",
        "    W, H,\n",
        "    per_pixel_laser_light,\n",
        "    laser_decoder_fn,\n",
        "    motorcycle_scene.z_proj_ground_truth,\n",
        "    motorcycle_scene.z_cam_ground_truth,\n",
        "    motorcycle_scene.rgb_ground_truth,\n",
        "    motorcycle_scene.proj_extrinsics,\n",
        "    motorcycle_scene.proj_intrinsics,\n",
        "    motorcycle_scene.cam_extrinsics,\n",
        "    motorcycle_scene.cam_intrinsics,\n",
        "    plot_skip = 100)\n",
        "\n",
        "# Display a progressive animation of the accumulated depth map\n",
        "depth_animation.generate()\n",
        "depth_animation.anim"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_Q7Uo3KM75Y6"
      },
      "source": [
        "You should see an animation of the depth recovery as we go through the set of structured light images and add to the depth map.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C2L_L8SNcb_k"
      },
      "source": [
        "## 2.3 Depth from Structured Light\n",
        "\n",
        "The laser light approach is elegant. It is easy to decode the matching pairs.\n",
        "\n",
        "However, you may have noticed that it is very *inefficient*. For a 1024x1024 image, you will have to produce over a **million** light images, making it  unviable in practice (unless you have a really high-power laser, but then you might set your scene on fire)\n",
        "\n",
        "In practice we use patterns of light or a spatially-varying light signal to obtain more corresponding pairs per projected image.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XU16HO3I5U0y"
      },
      "source": [
        "### 2.3. (a) Efficient Generator & Decoder Functions\n",
        "\n",
        "Your next (and final) task is to come up with a more efficient pair of methods.\n",
        "\n",
        "Fill in `my_structured_light_generator_fn()` and `my_decoder_fn()`. Remember that you can call `render_scene_with_projected_light()` to see what your lighting pattern looks like when applied to the scene.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "g5Ve08-tPCvu"
      },
      "outputs": [],
      "source": [
        "# @title 5 points\n",
        "def my_structured_light_generator_fn(W: int, H: int) -> Iterable[Tuple[torch.Tensor, Optional[Any]]]:\n",
        "    \"\"\"\n",
        "      Args:\n",
        "        W: width of the projector image in pixels\n",
        "        H: height of the projector image in pixels\n",
        "\n",
        "      Yields, 1 or more times, a tuple of:\n",
        "        light_image: image to be projected of shape (W, H)\n",
        "        extra_info: optional extra info to be passed onto the decoder fn.\n",
        "                    Use 'None' if you don't need this output.\n",
        "    \"\"\"\n",
        "\n",
        "    # Notes:\n",
        "    #\n",
        "    # This function should be a python 'generator'. Instead of using\n",
        "    # 'return', use 'yield' to return one or more images.\n",
        "    # This feature is handy since you can use a single pattern if you want, or an\n",
        "    # arbitrary number of different patterns (just remember to write your decoder\n",
        "    # accordingly).\n",
        "    # Alternatively, you can simply return a list of images (but generators are\n",
        "    # more memory-efficient)\n",
        "    #\n",
        "    # For our simple simulation setup, we have 'perfect' reflectance.\n",
        "    # That is, we recover the exact light value that we project, only the\n",
        "    # position on the image plane changes. This lets us recover depth using a\n",
        "    # single light image, i.e. using a single yield statement and without using\n",
        "    # the 'extra_info' argument.\n",
        "    #\n",
        "    # Hint: project a signal that is a function of uv_p, so that you can decode\n",
        "    # uv_p just by looking at the signal value at each uv_c.\n",
        "    #\n",
        "    # Alternatively, you can also use the provided stripes pattern, and 'scan' it\n",
        "    # horizontally while applying edge-detection to decode matching pairs.\n",
        "    # This approach is closer to what is used in real-world scenarios since we\n",
        "    # cannot assume a perfect reflectance model.\n",
        "    #\n",
        "\n",
        "    raise NotImplementedError\n",
        "\n",
        "\n",
        "def my_decoder_fn(\n",
        "    light_image: torch.Tensor,\n",
        "    camera_image: torch.Tensor,\n",
        "    extra_info: Any = None\n",
        "    ) -> Tuple[torch.Tensor, torch.Tensor]:\n",
        "    \"\"\"\n",
        "      Args:\n",
        "        light_image: projector's light image generated from the generator_fn\n",
        "        camera_image: camera image of the projected light\n",
        "        extra_info: optional extra information from the generator fn (See above)\n",
        "\n",
        "      Returns (a tuple of):\n",
        "        uv_p: list of pixel coordinates on the projector's image space with shape (N, 2)\n",
        "        uv_c: corresponding list of pixel coordinates on the camera's image space with shape (N, 2)\n",
        "\n",
        "      Note: N can be whatever number of pairs you are able to decode from a\n",
        "            single pair of structured light image and corresponding camera image.\n",
        "            The higher N is, the more efficient your structured light pattern.\n",
        "            You should recover at least 80% of pixels on this scene for full\n",
        "            credit (use the function below to verify this).\n",
        "            You might not need to use light_image or extra_info.\n",
        "    \"\"\"\n",
        "\n",
        "    raise NotImplementedError"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TIfXeEdXcP-c"
      },
      "outputs": [],
      "source": [
        "# verifies your generator and decoder functions behave correctly\n",
        "def check_generator_decoder_fns(\n",
        "      generator_fn,\n",
        "      decoder_fn,\n",
        "      scene,\n",
        "      max_relative_L1_error,\n",
        "      max_light_image_count,\n",
        "      min_recovered_fraction):\n",
        "\n",
        "    W, H = scene.z_cam_ground_truth.shape\n",
        "\n",
        "    n_light_images = len([image for image in generator_fn(W, H)])\n",
        "    recovered_depth, _ = recover_depth_map(\n",
        "                W, H,\n",
        "                generator_fn,\n",
        "                decoder_fn,\n",
        "                scene.z_proj_ground_truth,\n",
        "                scene.z_cam_ground_truth,\n",
        "                scene.rgb_ground_truth,\n",
        "                scene.proj_extrinsics,\n",
        "                scene.proj_intrinsics,\n",
        "                scene.cam_extrinsics,\n",
        "                scene.cam_intrinsics,\n",
        "                plot_skip = 1)\n",
        "    print(\"\")\n",
        "\n",
        "    non_nan_mask = ~recovered_depth.isnan()\n",
        "    relative_L1_error = (abs(recovered_depth[non_nan_mask] - scene.z_cam_ground_truth[non_nan_mask])\n",
        "                            / torch.max(scene.z_cam_ground_truth)).mean()\n",
        "    recovered_fraction = recovered_depth[non_nan_mask].numel() / scene.z_cam_ground_truth.numel()\n",
        "\n",
        "    correct = True\n",
        "    if relative_L1_error > max_relative_L1_error:\n",
        "        print(\"\\033[31m \", f'Expected relative L1 error below {max_relative_L1_error:.5f}. Got: {relative_L1_error:.5f}.', \" \\033[0m\")\n",
        "        correct = False\n",
        "\n",
        "    if n_light_images > max_light_image_count:\n",
        "        print(\"\\033[31m \", f'Expected fewer than {max_light_image_count} projector images. Got: {n_light_images}.', \" \\033[0m\")\n",
        "        correct = False\n",
        "\n",
        "    if recovered_fraction < min_recovered_fraction:\n",
        "        print(\"\\033[31m \", f'Expected atleast {(min_recovered_fraction * 100):.2f}% of depth image to be recovered. Got: {(recovered_fraction * 100):.2f}%.', \" \\033[0m\")\n",
        "        correct = False\n",
        "\n",
        "    if correct:\n",
        "        print(\"\\033[32m \", f'{generator_fn.__name__}/{decoder_fn.__name__}: Your generator-decoder pair works!', \" \\033[0m\")\n",
        "    else:\n",
        "        print(\"\\033[31m \", f'{generator_fn.__name__}/{decoder_fn.__name__}: Your generator-decoder does NOT work!', \" \\033[0m\")\n",
        "\n",
        "from math import sqrt\n",
        "\n",
        "# Together, the two functions are worth 5 points.\n",
        "check_generator_decoder_fns(\n",
        "    my_structured_light_generator_fn,\n",
        "    my_decoder_fn,\n",
        "    motorcycle_scene_large,\n",
        "    max_relative_L1_error=0.02,\n",
        "    max_light_image_count=(4 * sqrt(motorcycle_scene_large.z_cam_ground_truth.numel())),\n",
        "    min_recovered_fraction=0.8)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YZFN0TbDR5eF"
      },
      "source": [
        "Once you have an implementation, run the next cell to recover depth again, but with your new structured light method."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XTEn5H8aR5Gj"
      },
      "outputs": [],
      "source": [
        "W, H = motorcycle_scene_large.rgb_ground_truth.shape[:-1]\n",
        "\n",
        "# Recover a depth map using your structured light.\n",
        "# Together, the two functions are worth 5 points.\n",
        "z_recovered, depth_animation = recover_depth_map(\n",
        "    W, H,\n",
        "    my_structured_light_generator_fn,\n",
        "    my_decoder_fn,\n",
        "    motorcycle_scene_large.z_proj_ground_truth,\n",
        "    motorcycle_scene_large.z_cam_ground_truth,\n",
        "    motorcycle_scene_large.rgb_ground_truth,\n",
        "    motorcycle_scene_large.proj_extrinsics,\n",
        "    motorcycle_scene_large.proj_intrinsics,\n",
        "    motorcycle_scene_large.cam_extrinsics,\n",
        "    motorcycle_scene_large.cam_intrinsics,\n",
        "    plot_skip = 1)\n",
        "\n",
        "# Display a progressive animation of the accumulated depth map\n",
        "depth_animation.generate()\n",
        "depth_animation.anim"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eQuf9IrwSI79"
      },
      "source": [
        "You should be able to see a recovered depth map, as well a print statements showing how many light images you had to use.\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wG2PPOOCcfXH"
      },
      "source": [
        "You have completed Problem Set 6. See below for submission instructions\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YnoL7073RoiB"
      },
      "source": [
        "# Submission Instructions\n",
        "\n",
        "Be sure you have tested your code, and have not changed any function names. A minimal check that your code runs without error is to click \"Runtime -> Restart session and run all\". We will only consider regrade requests under exceptional circumstances.\n",
        "\n",
        "Use \"File -> Download -> Download `.ipynb`\" and upload the file to Gradescope. Do not submit a report or anything else: the autograder will give you 0 if you submit anything other than a .ipynb file."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XqCwQtcAcoTI"
      },
      "source": [
        "To learn more about modern inverse graphics techniques, consider taking _[6.S980 Machine Learning for Inverse Graphics](https://www.scenerepresentations.org/courses/inverse-graphics/)_ in the Fall.  "
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": [],
      "collapsed_sections": [
        "Q4dFIVuBrUch",
        "SZPPfdQ11TwP",
        "0xSnhuim2VZt",
        "Vq1rwQbY0dCV",
        "zcXctUr8MLWA",
        "Mggi7T2jUGyF",
        "eblV24_zI9PO",
        "w7RWEnooMOXJ",
        "fniw72sYzJBR",
        "wFDkiE4lwY8a"
      ],
      "toc_visible": true,
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}