{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copyright (c) Facebook, Inc. and its 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": [
    "If `torch`, `torchvision` and `pytorch3d` are not installed, run the following cell:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import matplotlib.pyplot as plt\n",
    "from skimage.io import imread\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.renderer import (\n",
    "    look_at_view_transform,\n",
    "    FoVOrthographicCameras, \n",
    "    PointsRasterizationSettings,\n",
    "    PointsRenderer,\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** (orthgraphic/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",
    "renderer = PointsRenderer(\n",
    "    rasterizer=PointsRasterizer(cameras=cameras, raster_settings=raster_settings),\n",
    "    compositor=AlphaCompositor()\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.grid(\"off\")\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.grid(\"off\")\n",
    "plt.axis(\"off\");"
   ]
  }
 ],
 "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": "pytorch3d (local)",
   "language": "python",
   "name": "pytorch3d_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": 2
}
