docs/tutorials/render_colored_points.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copyright (c) Meta Platforms, Inc. and affiliates. All rights reserved."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Render a colored point cloud\n",
    "\n",
    "This tutorial shows how to:\n",
    "- set up a renderer \n",
    "- render the point cloud \n",
    "- vary the rendering settings such as compositing and camera position"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import modules"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Ensure `torch` and `torchvision` are installed. If `pytorch3d` is not installed, install it using the following cell:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import torch\n",
    "import subprocess\n",
    "need_pytorch3d=False\n",
    "try:\n",
    "    import pytorch3d\n",
    "except ModuleNotFoundError:\n",
    "    need_pytorch3d=True\n",
    "if need_pytorch3d:\n",
    "    pyt_version_str=torch.__version__.split(\"+\")[0].replace(\".\", \"\")\n",
    "    version_str=\"\".join([\n",
    "        f\"py3{sys.version_info.minor}_cu\",\n",
    "        torch.version.cuda.replace(\".\",\"\"),\n",
    "        f\"_pyt{pyt_version_str}\"\n",
    "    ])\n",
    "    !pip install iopath\n",
    "    if sys.platform.startswith(\"linux\"):\n",
    "        print(\"Trying to install wheel for PyTorch3D\")\n",
    "        !pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/{version_str}/download.html\n",
    "        pip_list = !pip freeze\n",
    "        need_pytorch3d = not any(i.startswith(\"pytorch3d==\") for  i in pip_list)\n",
    "    if need_pytorch3d:\n",
    "        print(f\"failed to find/install wheel for {version_str}\")\n",
    "if need_pytorch3d:\n",
    "    print(\"Installing PyTorch3D from source\")\n",
    "    !pip install ninja\n",
    "    !pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Util function for loading point clouds|\n",
    "import numpy as np\n",
    "\n",
    "# Data structures and functions for rendering\n",
    "from pytorch3d.structures import Pointclouds\n",
    "from pytorch3d.vis.plotly_vis import AxisArgs, plot_batch_individually, plot_scene\n",
    "from pytorch3d.renderer import (\n",
    "    look_at_view_transform,\n",
    "    FoVOrthographicCameras, \n",
    "    PointsRasterizationSettings,\n",
    "    PointsRenderer,\n",
    "    PulsarPointsRenderer,\n",
    "    PointsRasterizer,\n",
    "    AlphaCompositor,\n",
    "    NormWeightedCompositor\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load a point cloud and corresponding colors\n",
    "\n",
    "Load and create a **Point Cloud** object. \n",
    "\n",
    "**Pointclouds** is a unique datastructure provided in PyTorch3D for working with batches of point clouds of different sizes. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If running this notebook using **Google Colab**, run the following cell to fetch the pointcloud data and save it at the path `data/PittsburghBridge`:\n",
    "If running locally, the data is already available at the correct path. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p data/PittsburghBridge\n",
    "!wget -P data/PittsburghBridge https://dl.fbaipublicfiles.com/pytorch3d/data/PittsburghBridge/pointcloud.npz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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, \"PittsburghBridge/pointcloud.npz\")\n",
    "\n",
    "# Load point cloud\n",
    "pointcloud = np.load(obj_filename)\n",
    "verts = torch.Tensor(pointcloud['verts']).to(device)\n",
    "        \n",
    "rgb = torch.Tensor(pointcloud['rgb']).to(device)\n",
    "\n",
    "point_cloud = Pointclouds(points=[verts], features=[rgb])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 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 an **orthographic camera**, and applies **alpha compositing**. Then we learn how to vary different components using the modular API.  \n",
    "\n",
    "[1] <a href=\"https://arxiv.org/abs/1912.08804\">SynSin: End to end View Synthesis from a Single Image.</a> Olivia Wiles, Georgia Gkioxari, Richard Szeliski, Justin Johnson. CVPR 2020."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize a camera.\n",
    "R, T = look_at_view_transform(20, 10, 0)\n",
    "cameras = FoVOrthographicCameras(device=device, R=R, T=T, znear=0.01)\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. Refer to raster_points.py for explanations of these parameters. \n",
    "raster_settings = PointsRasterizationSettings(\n",
    "    image_size=512, \n",
    "    radius = 0.003,\n",
    "    points_per_pixel = 10\n",
    ")\n",
    "\n",
    "\n",
    "# Create a points renderer by compositing points using an alpha compositor (nearer points\n",
    "# are weighted more heavily). See [1] for an explanation.\n",
    "rasterizer = PointsRasterizer(cameras=cameras, raster_settings=raster_settings)\n",
    "renderer = PointsRenderer(\n",
    "    rasterizer=rasterizer,\n",
    "    compositor=AlphaCompositor()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "images = renderer(point_cloud)\n",
    "plt.figure(figsize=(10, 10))\n",
    "plt.imshow(images[0, ..., :3].cpu().numpy())\n",
    "plt.axis(\"off\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will now modify the **renderer** to use **alpha compositing** with a set background color. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "renderer = PointsRenderer(\n",
    "    rasterizer=rasterizer,\n",
    "    # Pass in background_color to the alpha compositor, setting the background color \n",
    "    # to the 3 item tuple, representing rgb on a scale of 0 -> 1, in this case blue\n",
    "    compositor=AlphaCompositor(background_color=(0, 0, 1))\n",
    ")\n",
    "images = renderer(point_cloud)\n",
    "\n",
    "plt.figure(figsize=(10, 10))\n",
    "plt.imshow(images[0, ..., :3].cpu().numpy())\n",
    "plt.axis(\"off\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this example we will first create a **renderer** which uses an **orthographic camera**, and applies **weighted compositing**. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize a camera.\n",
    "R, T = look_at_view_transform(20, 10, 0)\n",
    "cameras = FoVOrthographicCameras(device=device, R=R, T=T, znear=0.01)\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. Refer to rasterize_points.py for explanations of these parameters. \n",
    "raster_settings = PointsRasterizationSettings(\n",
    "    image_size=512, \n",
    "    radius = 0.003,\n",
    "    points_per_pixel = 10\n",
    ")\n",
    "\n",
    "\n",
    "# Create a points renderer by compositing points using an weighted compositor (3D points are\n",
    "# weighted according to their distance to a pixel and accumulated using a weighted sum)\n",
    "renderer = PointsRenderer(\n",
    "    rasterizer=PointsRasterizer(cameras=cameras, raster_settings=raster_settings),\n",
    "    compositor=NormWeightedCompositor()\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "images = renderer(point_cloud)\n",
    "plt.figure(figsize=(10, 10))\n",
    "plt.imshow(images[0, ..., :3].cpu().numpy())\n",
    "plt.axis(\"off\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will now modify the **renderer** to use **weighted compositing** with a set background color. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "renderer = PointsRenderer(\n",
    "    rasterizer=PointsRasterizer(cameras=cameras, raster_settings=raster_settings),\n",
    "    # Pass in background_color to the norm weighted compositor, setting the background color \n",
    "    # to the 3 item tuple, representing rgb on a scale of 0 -> 1, in this case red\n",
    "    compositor=NormWeightedCompositor(background_color=(1,0,0))\n",
    ")\n",
    "images = renderer(point_cloud)\n",
    "plt.figure(figsize=(10, 10))\n",
    "plt.imshow(images[0, ..., :3].cpu().numpy())\n",
    "plt.axis(\"off\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using the pulsar backend\n",
    "\n",
    "Switching to the pulsar backend is easy! The pulsar backend has a compositor built-in, so the `compositor` argument is not required when creating it (a warning will be displayed if you provide it nevertheless). It pre-allocates memory on the rendering device, that's why it needs the `n_channels` at construction time.\n",
    "\n",
    "All parameters for the renderer forward function are batch-wise except the background color (in this example, `gamma`) and you have to provide as many values as you have examples in your batch. The background color is optional and by default set to all zeros. You can find a detailed explanation of how gamma influences the rendering function here in the paper [Fast Differentiable Raycasting for Neural Rendering using\n",
    "Sphere-based Representations](https://arxiv.org/pdf/2004.07484.pdf).\n",
    "\n",
    "You can also use the `native` backend for the pulsar backend which already provides access to point opacity. The native backend can be imported from `pytorch3d.renderer.points.pulsar`; you can find examples for this in the folder `docs/examples`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "renderer = PulsarPointsRenderer(\n",
    "    rasterizer=PointsRasterizer(cameras=cameras, raster_settings=raster_settings),\n",
    "    n_channels=4\n",
    ").to(device)\n",
    "\n",
    "images = renderer(point_cloud, gamma=(1e-4,),\n",
    "                  bg_col=torch.tensor([0.0, 1.0, 0.0, 1.0], dtype=torch.float32, device=device))\n",
    "plt.figure(figsize=(10, 10))\n",
    "plt.imshow(images[0, ..., :3].cpu().numpy())\n",
    "plt.axis(\"off\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### View pointclouds in Plotly figures\n",
    "\n",
    "Here we use the PyTorch3D function `plot_scene` to render the pointcloud in a Plotly figure. `plot_scene` returns a plotly figure with trace and subplots defined by the input."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_scene({\n",
    "    \"Pointcloud\": {\n",
    "        \"person\": point_cloud\n",
    "    }\n",
    "})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will now render a batch of pointclouds. The first pointcloud is the same as above, and the second is all-black and offset by 2 in all dimensions so we can see them on the same plot. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "point_cloud_batch = Pointclouds(points=[verts, verts + 2], features=[rgb, torch.zeros_like(rgb)])\n",
    "# render both in the same plot in different traces\n",
    "fig = plot_scene({\n",
    "    \"Pointcloud\": {\n",
    "        \"person\": point_cloud_batch[0],\n",
    "        \"person2\": point_cloud_batch[1]\n",
    "    }\n",
    "})\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# render both in the same plot in one trace\n",
    "fig = plot_scene({\n",
    "    \"Pointcloud\": {\n",
    "        \"2 people\": point_cloud_batch\n",
    "    }\n",
    "})\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For batches, we can also use `plot_batch_individually` to avoid constructing the scene dictionary ourselves."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# render both in 1 row in different subplots\n",
    "fig2 = plot_batch_individually(point_cloud_batch, ncols=2)\n",
    "fig2.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# modify the plotly figure height and width\n",
    "fig2.update_layout(height=500, width=500)\n",
    "fig2.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also modify the axis arguments and axis backgrounds for either function, and title our plots in `plot_batch_individually`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig3 = plot_batch_individually(\n",
    "    point_cloud_batch, \n",
    "    xaxis={\"backgroundcolor\":\"rgb(200, 200, 230)\"},\n",
    "    yaxis={\"backgroundcolor\":\"rgb(230, 200, 200)\"},\n",
    "    zaxis={\"backgroundcolor\":\"rgb(200, 230, 200)\"}, \n",
    "    subplot_titles=[\"Pointcloud1\", \"Pointcloud2\"], # this should have a title for each subplot, titles can be \"\"\n",
    "    axis_args=AxisArgs(showgrid=True))\n",
    "fig3.show()"
   ]
  }
 ],
 "metadata": {
  "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
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}