render_textured_meshes.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "_Ip8kp4TfBLZ"
   },
   "outputs": [],
   "source": [
    "# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "kuXHJv44fBLe"
   },
   "source": [
    "# Render a textured mesh\n",
    "\n",
    "This tutorial shows how to:\n",
    "- load a mesh and textures from an `.obj` file. \n",
    "- set up a renderer \n",
    "- render the mesh \n",
    "- vary the rendering settings such as lighting and camera position\n",
    "- use the batching features of the pytorch3d API to render the mesh from different viewpoints"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Bnj3THhzfBLf"
   },
   "source": [
    "## 0. Install and Import modules"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "okLalbR_g7NS"
   },
   "source": [
    "If `torch`, `torchvision` and `pytorch3d` are not installed, run the following cell:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 717
    },
    "colab_type": "code",
    "id": "musUWTglgxSB",
    "outputId": "16d1a1b2-3f7f-43ed-ca28-a4d236cc0572"
   },
   "outputs": [],
   "source": [
    "!pip install torch torchvision\n",
    "import sys\n",
    "import torch\n",
    "if torch.__version__=='1.6.0+cu101' and sys.platform.startswith('linux'):\n",
    "    !pip install pytorch3d\n",
    "else:\n",
    "    !pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "nX99zdoffBLg"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import matplotlib.pyplot as plt\n",
    "from skimage.io import imread\n",
    "\n",
    "# Util function for loading meshes\n",
    "from pytorch3d.io import load_objs_as_meshes, load_obj\n",
    "\n",
    "# Data structures and functions for rendering\n",
    "from pytorch3d.structures import Meshes\n",
    "from pytorch3d.vis import AxisArgs, plot_meshes\n",
    "from pytorch3d.renderer import (\n",
    "    look_at_view_transform,\n",
    "    FoVPerspectiveCameras, \n",
    "    PointLights, \n",
    "    DirectionalLights, \n",
    "    Materials, \n",
    "    RasterizationSettings, \n",
    "    MeshRenderer, \n",
    "    MeshRasterizer,  \n",
    "    SoftPhongShader,\n",
    "    TexturesUV,\n",
    "    TexturesVertex\n",
    ")\n",
    "\n",
    "# add path for demo utils functions \n",
    "import sys\n",
    "import os\n",
    "sys.path.append(os.path.abspath(''))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Lxmehq6Zhrzv"
   },
   "source": [
    "If using **Google Colab**, fetch the utils file for plotting image grids:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 204
    },
    "colab_type": "code",
    "id": "HZozr3Pmho-5",
    "outputId": "be5eb60d-5f65-4db1-cca0-44ee68c8f5fd"
   },
   "outputs": [],
   "source": [
    "!wget https://raw.githubusercontent.com/facebookresearch/pytorch3d/master/docs/tutorials/utils/plot_image_grid.py\n",
    "from plot_image_grid import image_grid"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "g4B62MzYiJUM"
   },
   "source": [
    "OR if running **locally** uncomment and run the following cell:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "paJ4Im8ahl7O"
   },
   "outputs": [],
   "source": [
    "# from utils import image_grid"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "5jGq772XfBLk"
   },
   "source": [
    "### 1. Load a mesh and texture file\n",
    "\n",
    "Load an `.obj` file and it's associated `.mtl` file and create a **Textures** and **Meshes** object. \n",
    "\n",
    "**Meshes** is a unique datastructure provided in PyTorch3D for working with batches of meshes of different sizes. \n",
    "\n",
    "**TexturesUV** is an auxillary datastructure for storing vertex uv and texture maps for meshes. \n",
    "\n",
    "**Meshes** has several class methods which are used throughout the rendering pipeline."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "a8eU4zo5jd_H"
   },
   "source": [
    "If running this notebook using **Google Colab**, run the following cell to fetch the mesh obj and texture files and save it at the path `data/cow_mesh`:\n",
    "If running locally, the data is already available at the correct path. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 578
    },
    "colab_type": "code",
    "id": "tTm0cVuOjb1W",
    "outputId": "6cd7e2ec-65e1-4dcc-99e8-c347bc504f0a"
   },
   "outputs": [],
   "source": [
    "!mkdir -p data/cow_mesh\n",
    "!wget -P data/cow_mesh https://dl.fbaipublicfiles.com/pytorch3d/data/cow_mesh/cow.obj\n",
    "!wget -P data/cow_mesh https://dl.fbaipublicfiles.com/pytorch3d/data/cow_mesh/cow.mtl\n",
    "!wget -P data/cow_mesh https://dl.fbaipublicfiles.com/pytorch3d/data/cow_mesh/cow_texture.png"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "gi5Kd0GafBLl"
   },
   "outputs": [],
   "source": [
    "# Setup\n",
    "if torch.cuda.is_available():\n",
    "    device = torch.device(\"cuda:0\")\n",
    "    torch.cuda.set_device(device)\n",
    "else:\n",
    "    device = torch.device(\"cpu\")\n",
    "\n",
    "# Set paths\n",
    "DATA_DIR = \"./data\"\n",
    "obj_filename = os.path.join(DATA_DIR, \"cow_mesh/cow.obj\")\n",
    "\n",
    "# Load obj file\n",
    "mesh = load_objs_as_meshes([obj_filename], device=device)\n",
    "texture_image=mesh.textures.maps_padded()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "5APAQs6-fBLp"
   },
   "source": [
    "#### Let's visualize the texture map"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 428
    },
    "colab_type": "code",
    "id": "YipUhrIHfBLq",
    "outputId": "48987b1d-5cc1-4c2a-cb3c-713d64f6a38d"
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(7,7))\n",
    "plt.imshow(texture_image.squeeze().cpu().numpy())\n",
    "plt.grid(\"off\");\n",
    "plt.axis('off');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "GcnG6XJ6fBLu"
   },
   "source": [
    "## 2. Create a renderer\n",
    "\n",
    "A renderer in PyTorch3D is composed of a **rasterizer** and a **shader** which each have a number of subcomponents such as a **camera** (orthographic/perspective). Here we initialize some of these components and use default values for the rest.\n",
    "\n",
    "In this example we will first create a **renderer** which uses a **perspective camera**, a **point light** and applies **phong shading**. Then we learn how to vary different components using the modular API.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "dX466mWnfBLv"
   },
   "outputs": [],
   "source": [
    "# Initialize a camera.\n",
    "# With world coordinates +Y up, +X left and +Z in, the front of the cow is facing the -Z direction. \n",
    "# So we move the camera by 180 in the azimuth direction so it is facing the front of the cow. \n",
    "R, T = look_at_view_transform(2.7, 0, 180) \n",
    "cameras = FoVPerspectiveCameras(device=device, R=R, T=T)\n",
    "\n",
    "# Define the settings for rasterization and shading. Here we set the output image to be of size\n",
    "# 512x512. As we are rendering images for visualization purposes only we will set faces_per_pixel=1\n",
    "# and blur_radius=0.0. We also set bin_size and max_faces_per_bin to None which ensure that \n",
    "# the faster coarse-to-fine rasterization method is used. Refer to rasterize_meshes.py for \n",
    "# explanations of these parameters. Refer to docs/notes/renderer.md for an explanation of \n",
    "# the difference between naive and coarse-to-fine rasterization. \n",
    "raster_settings = RasterizationSettings(\n",
    "    image_size=512, \n",
    "    blur_radius=0.0, \n",
    "    faces_per_pixel=1, \n",
    ")\n",
    "\n",
    "# Place a point light in front of the object. As mentioned above, the front of the cow is facing the \n",
    "# -z direction. \n",
    "lights = PointLights(device=device, location=[[0.0, 0.0, -3.0]])\n",
    "\n",
    "# Create a phong renderer by composing a rasterizer and a shader. The textured phong shader will \n",
    "# interpolate the texture uv coordinates for each vertex, sample from a texture image and \n",
    "# apply the Phong lighting model\n",
    "renderer = MeshRenderer(\n",
    "    rasterizer=MeshRasterizer(\n",
    "        cameras=cameras, \n",
    "        raster_settings=raster_settings\n",
    "    ),\n",
    "    shader=SoftPhongShader(\n",
    "        device=device, \n",
    "        cameras=cameras,\n",
    "        lights=lights\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "KyOY5qXvfBLz"
   },
   "source": [
    "## 3. Render the mesh"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "8VkRA4qJfBL0"
   },
   "source": [
    "The light is in front of the object so it is bright and the image has specular highlights."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 592
    },
    "colab_type": "code",
    "id": "gBLZH8iUfBL1",
    "outputId": "cc3cd3f0-189e-4497-ce47-e64b4da542e8"
   },
   "outputs": [],
   "source": [
    "images = renderer(mesh)\n",
    "plt.figure(figsize=(10, 10))\n",
    "plt.imshow(images[0, ..., :3].cpu().numpy())\n",
    "plt.grid(\"off\");\n",
    "plt.axis(\"off\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "k161XF3sfBL5"
   },
   "source": [
    "## 4. Move the light behind the object and re-render\n",
    "\n",
    "We can pass arbirary keyword arguments to the `rasterizer`/`shader` via the call to the `renderer` so the renderer does not need to be reinitialized if any of the settings change/\n",
    "\n",
    "In this case, we can simply update the location of the lights and pass them into the call to the renderer. \n",
    "\n",
    "The image is now dark as there is only ambient lighting, and there are no specular highlights."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "BdWkkeibfBL6"
   },
   "outputs": [],
   "source": [
    "# Now move the light so it is on the +Z axis which will be behind the cow. \n",
    "lights.location = torch.tensor([0.0, 0.0, +1.0], device=device)[None]\n",
    "images = renderer(mesh, lights=lights)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 592
    },
    "colab_type": "code",
    "id": "UmV3j1YffBL9",
    "outputId": "2e8edca0-5bd8-4a2f-a160-83c4b0520123"
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(10, 10))\n",
    "plt.imshow(images[0, ..., :3].cpu().numpy())\n",
    "plt.grid(\"off\");\n",
    "plt.axis(\"off\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "t93aVotMfBMB"
   },
   "source": [
    "## 5. Rotate the object, modify the material properties or light properties\n",
    "\n",
    "We can also change many other settings in the rendering pipeline. Here we:\n",
    "\n",
    "- change the **viewing angle** of the camera\n",
    "- change the **position** of the point light\n",
    "- change the **material reflectance** properties of the mesh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "4mYXYziefBMB"
   },
   "outputs": [],
   "source": [
    "# Rotate the object by increasing the elevation and azimuth angles\n",
    "R, T = look_at_view_transform(dist=2.7, elev=10, azim=-150)\n",
    "cameras = FoVPerspectiveCameras(device=device, R=R, T=T)\n",
    "\n",
    "# Move the light location so the light is shining on the cow's face.  \n",
    "lights.location = torch.tensor([[2.0, 2.0, -2.0]], device=device)\n",
    "\n",
    "# Change specular color to green and change material shininess \n",
    "materials = Materials(\n",
    "    device=device,\n",
    "    specular_color=[[0.0, 1.0, 0.0]],\n",
    "    shininess=10.0\n",
    ")\n",
    "\n",
    "# Re render the mesh, passing in keyword arguments for the modified components.\n",
    "images = renderer(mesh, lights=lights, materials=materials, cameras=cameras)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 592
    },
    "colab_type": "code",
    "id": "rHIxIfh5fBME",
    "outputId": "1ca2d337-2983-478f-b3c9-d64b84ba1a31"
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(10, 10))\n",
    "plt.imshow(images[0, ..., :3].cpu().numpy())\n",
    "plt.grid(\"off\");\n",
    "plt.axis(\"off\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "17c4xmtyfBMH"
   },
   "source": [
    "## 6. Batched Rendering\n",
    "\n",
    "One of the core design choices of the PyTorch3D API is to support **batched inputs for all components**. \n",
    "The renderer and associated components can take batched inputs and **render a batch of output images in one forward pass**. We will now use this feature to render the mesh from many different viewpoints.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "CDQKebNNfBMI"
   },
   "outputs": [],
   "source": [
    "# Set batch size - this is the number of different viewpoints from which we want to render the mesh.\n",
    "batch_size = 20\n",
    "\n",
    "# Create a batch of meshes by repeating the cow mesh and associated textures. \n",
    "# Meshes has a useful `extend` method which allows us do this very easily. \n",
    "# This also extends the textures. \n",
    "meshes = mesh.extend(batch_size)\n",
    "\n",
    "# Get a batch of viewing angles. \n",
    "elev = torch.linspace(0, 180, batch_size)\n",
    "azim = torch.linspace(-180, 180, batch_size)\n",
    "\n",
    "# All the cameras helper methods support mixed type inputs and broadcasting. So we can \n",
    "# view the camera from the same distance and specify dist=2.7 as a float,\n",
    "# and then specify elevation and azimuth angles for each viewpoint as tensors. \n",
    "R, T = look_at_view_transform(dist=2.7, elev=elev, azim=azim)\n",
    "cameras = FoVPerspectiveCameras(device=device, R=R, T=T)\n",
    "\n",
    "# Move the light back in front of the cow which is facing the -z direction.\n",
    "lights.location = torch.tensor([[0.0, 0.0, -3.0]], device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "gyYJCwEDfBML"
   },
   "outputs": [],
   "source": [
    "# We can pass arbirary keyword arguments to the rasterizer/shader via the renderer\n",
    "# so the renderer does not need to be reinitialized if any of the settings change.\n",
    "images = renderer(meshes, cameras=cameras, lights=lights)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "image_grid(images.cpu().numpy(), rows=4, cols=5, rgb=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Plotly visualization \n",
    "If you only want to visualize a mesh, you don't really need to use a differentiable renderer - instead we support plotting of Meshes with plotly. For these Meshes, we use TexturesVertex to define a texture for the rendering.\n",
    "`plot_meshes` creates a Plotly figure with a trace for each Meshes object. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "verts, faces_idx, _ = load_obj(obj_filename)\n",
    "faces = faces_idx.verts_idx\n",
    "\n",
    "# Initialize each vertex to be white in color.\n",
    "verts_rgb = torch.ones_like(verts)[None]  # (1, V, 3)\n",
    "textures = TexturesVertex(verts_features=verts_rgb.to(device))\n",
    "\n",
    "# Create a Meshes object\n",
    "mesh = Meshes(\n",
    "    verts=[verts.to(device)],   \n",
    "    faces=[faces.to(device)],\n",
    "    textures=textures\n",
    ")\n",
    "\n",
    "# Render the plotly figure\n",
    "fig = plot_meshes(mesh)\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# use Plotly's default colors (no texture)\n",
    "mesh = Meshes(\n",
    "    verts=[verts.to(device)],   \n",
    "    faces=[faces.to(device)]\n",
    ")\n",
    "\n",
    "# Render the plotly figure\n",
    "plot_meshes(mesh)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a batch of meshes, and offset one to prevent overlap\n",
    "mesh_batch = Meshes(\n",
    "    verts=[verts.to(device), (verts + 2).to(device)],   \n",
    "    faces=[faces.to(device), faces.to(device)]\n",
    ")\n",
    "\n",
    "plot_meshes(mesh_batch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# extend the batch to have 4 meshes\n",
    "mesh = mesh_batch.extend(2)\n",
    "\n",
    "# visualize the batch in different subplots, 2 per row\n",
    "# title the subplots\n",
    "fig = plot_meshes(mesh, in_subplots=True, ncols=2,\n",
    "           subplot_titles = [\"cow\" + str(i) for i in range(1,5)]  # this should have a title for each subplot, titles can be \"\"\n",
    "           )\n",
    "# we can update the figure height and width\n",
    "fig.update_layout(height=500, width=500)\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also modify the axis arguments and axis backgrounds."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig2 = plot_meshes(mesh_batch, xaxis={\"backgroundcolor\":\"rgb(200, 200, 230)\"},\n",
    "                        yaxis={\"backgroundcolor\":\"rgb(230, 200, 200)\"},\n",
    "                        zaxis={\"backgroundcolor\":\"rgb(200, 230, 200)\"}, \n",
    "                        subplot_titles=[\"2 meshes\"],\n",
    "                        axis_args=AxisArgs(showgrid=True))\n",
    "fig2.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "t3qphI1ElUb5"
   },
   "source": [
    "## 8. Conclusion\n",
    "In this tutorial we learnt how to **load** a textured mesh from an obj file, initialize a PyTorch3D datastructure called **Meshes**, set up an **Renderer** consisting of a **Rasterizer** and a **Shader**, and modify several components of the rendering pipeline. We also learned how to render Meshes in Plotly figures."
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "anp_metadata": {
   "path": "notebooks/render_textured_meshes.ipynb"
  },
  "bento_stylesheets": {
   "bento/extensions/flow/main.css": true,
   "bento/extensions/kernel_selector/main.css": true,
   "bento/extensions/kernel_ui/main.css": true,
   "bento/extensions/new_kernel/main.css": true,
   "bento/extensions/system_usage/main.css": true,
   "bento/extensions/theme/main.css": true
  },
  "colab": {
   "name": "render_textured_meshes.ipynb",
   "provenance": []
  },
  "disseminate_notebook_info": {
   "backup_notebook_id": "569222367081034"
  },
  "kernelspec": {
   "display_name": "pytorch3d_etc (local)",
   "language": "python",
   "name": "pytorch3d_etc_local"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.5+"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}