docs/tutorials/render_densepose.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 DensePose \n",
    "\n",
    "DensePose refers to dense human pose representation: https://github.com/facebookresearch/DensePose. \n",
    "In this tutorial, we provide an example of using DensePose data in PyTorch3D.\n",
    "\n",
    "This tutorial shows how to:\n",
    "- load a mesh and textures from densepose `.mat` and `.pkl` files\n",
    "- set up a renderer \n",
    "- render the mesh \n",
    "- vary the rendering settings such as lighting and camera position"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Bnj3THhzfBLf"
   },
   "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": [
    "# We also install chumpy as it is needed to load the SMPL model pickle file.\n",
    "!pip install chumpy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "# libraries for reading data from files\n",
    "from scipy.io import loadmat\n",
    "from PIL import Image\n",
    "import pickle\n",
    "\n",
    "# Data structures and functions for rendering\n",
    "from pytorch3d.structures import 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",
    ")\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": {},
   "source": [
    "## Load the SMPL model\n",
    "\n",
    "#### Download the SMPL model\n",
    "- Go to https://smpl.is.tue.mpg.de/download.php and sign up.\n",
    "- Download SMPL for Python Users and unzip.\n",
    "- Copy the file male template file **'models/basicModel_m_lbs_10_207_0_v1.0.0.pkl'** to the data/DensePose/ folder.\n",
    "   - rename the file to **'smpl_model.pkl'** or rename the string where it's commented below\n",
    "   \n",
    "If running this notebook using Google Colab, run the following cell to fetch the texture and UV values and save it at the correct path."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Texture image\n",
    "!wget -P data/DensePose https://raw.githubusercontent.com/facebookresearch/DensePose/master/DensePoseData/demo_data/texture_from_SURREAL.png\n",
    "\n",
    "# UV_processed.mat\n",
    "!wget https://dl.fbaipublicfiles.com/densepose/densepose_uv_data.tar.gz\n",
    "!tar xvf densepose_uv_data.tar.gz -C data/DensePose\n",
    "!rm densepose_uv_data.tar.gz"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load our texture UV data and our SMPL data, with some processing to correct data values and format."
   ]
  },
  {
   "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",
    "data_filename = os.path.join(DATA_DIR, \"DensePose/UV_Processed.mat\")\n",
    "tex_filename = os.path.join(DATA_DIR,\"DensePose/texture_from_SURREAL.png\")\n",
    "# rename your .pkl file or change this string\n",
    "verts_filename = os.path.join(DATA_DIR, \"DensePose/smpl_model.pkl\")\n",
    "\n",
    "\n",
    "# Load SMPL and texture data\n",
    "with open(verts_filename, 'rb') as f:\n",
    "    data = pickle.load(f, encoding='latin1') \n",
    "    v_template = torch.Tensor(data['v_template']).to(device) # (6890, 3)\n",
    "ALP_UV = loadmat(data_filename)\n",
    "with Image.open(tex_filename) as image:\n",
    "    np_image = np.asarray(image.convert(\"RGB\")).astype(np.float32)\n",
    "tex = torch.from_numpy(np_image / 255.)[None].to(device)\n",
    "\n",
    "verts = torch.from_numpy((ALP_UV[\"All_vertices\"]).astype(int)).squeeze().to(device) # (7829,)\n",
    "U = torch.Tensor(ALP_UV['All_U_norm']).to(device) # (7829, 1)\n",
    "V = torch.Tensor(ALP_UV['All_V_norm']).to(device) # (7829, 1)\n",
    "faces = torch.from_numpy((ALP_UV['All_Faces'] - 1).astype(int)).to(device)  # (13774, 3)\n",
    "face_indices = torch.Tensor(ALP_UV['All_FaceIndices']).squeeze()  # (13774,)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Display the texture image\n",
    "plt.figure(figsize=(10, 10))\n",
    "plt.imshow(tex.squeeze(0).cpu())\n",
    "plt.axis(\"off\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In DensePose, the body mesh is split into 24 parts. In the texture image, we can see the 24 parts are separated out into individual (200, 200) images per body part.  The convention in DensePose is that each face in the mesh is associated with a body part (given by the face_indices tensor above). The vertex UV values (in the range [0, 1]) for each face are specific to the (200, 200) size texture map for the part of the body that the mesh face corresponds to. We cannot use them directly with the entire texture map. We have to offset the vertex UV values depending on what body part the associated face corresponds to."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Map each face to a (u, v) offset\n",
    "offset_per_part = {}\n",
    "already_offset = set()\n",
    "cols, rows = 4, 6\n",
    "for i, u in enumerate(np.linspace(0, 1, cols, endpoint=False)):\n",
    "    for j, v in enumerate(np.linspace(0, 1, rows, endpoint=False)):\n",
    "        part = rows * i + j + 1  # parts are 1-indexed in face_indices\n",
    "        offset_per_part[part] = (u, v)\n",
    "\n",
    "U_norm = U.clone()\n",
    "V_norm = V.clone()\n",
    "\n",
    "# iterate over faces and offset the corresponding vertex u and v values\n",
    "for i in range(len(faces)):\n",
    "    face_vert_idxs = faces[i]\n",
    "    part = face_indices[i]\n",
    "    offset_u, offset_v = offset_per_part[int(part.item())]\n",
    "    \n",
    "    for vert_idx in face_vert_idxs:   \n",
    "        # vertices are reused, but we don't want to offset multiple times\n",
    "        if vert_idx.item() not in already_offset:\n",
    "            # offset u value\n",
    "            U_norm[vert_idx] = U[vert_idx] / cols + offset_u\n",
    "            # offset v value\n",
    "            # this also flips each part locally, as each part is upside down\n",
    "            V_norm[vert_idx] = (1 - V[vert_idx]) / rows + offset_v\n",
    "            # add vertex to our set tracking offsetted vertices\n",
    "            already_offset.add(vert_idx.item())\n",
    "\n",
    "# invert V values\n",
    "V_norm = 1 - V_norm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create our verts_uv values\n",
    "verts_uv = torch.cat([U_norm[None],V_norm[None]], dim=2) # (1, 7829, 2)\n",
    "\n",
    "# There are 6890 xyz vertex coordinates but 7829 vertex uv coordinates. \n",
    "# This is because the same vertex can be shared by multiple faces where each face may correspond to a different body part.  \n",
    "# Therefore when initializing the Meshes class,\n",
    "# we need to map each of the vertices referenced by the DensePose faces (in verts, which is the \"All_vertices\" field)\n",
    "# to the correct xyz coordinate in the SMPL template mesh.\n",
    "v_template_extended = v_template[verts-1][None] # (1, 7829, 3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create our textured mesh \n",
    "\n",
    "**Meshes** is a unique datastructure provided in PyTorch3D for working with batches of meshes of different sizes.\n",
    "\n",
    "**TexturesUV** is an auxiliary datastructure for storing vertex uv and texture maps for meshes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "texture = TexturesUV(maps=tex, faces_uvs=faces[None], verts_uvs=verts_uv)\n",
    "mesh = Meshes(v_template_extended, faces[None], texture)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create a renderer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize a camera.\n",
    "# World coordinates +Y up, +X left and +Z in.\n",
    "R, T = look_at_view_transform(2.7, 0, 0) \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. \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 person. \n",
    "lights = PointLights(device=device, location=[[0.0, 0.0, 2.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": {},
   "source": [
    "Render the textured mesh we created from the SMPL model and texture map."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "images = renderer(mesh)\n",
    "plt.figure(figsize=(10, 10))\n",
    "plt.imshow(images[0, ..., :3].cpu().numpy())\n",
    "plt.axis(\"off\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Different view and lighting of the body\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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Rotate the person by increasing the elevation and azimuth angles to view the back of the person from above. \n",
    "R, T = look_at_view_transform(2.7, 10, 180)\n",
    "cameras = FoVPerspectiveCameras(device=device, R=R, T=T)\n",
    "\n",
    "# Move the light location so the light is shining on the person's back.  \n",
    "lights.location = torch.tensor([[2.0, 2.0, -2.0]], device=device)\n",
    "\n",
    "# Re render the mesh, passing in keyword arguments for the modified components.\n",
    "images = renderer(mesh, lights=lights, cameras=cameras)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(10, 10))\n",
    "plt.imshow(images[0, ..., :3].cpu().numpy())\n",
    "plt.axis(\"off\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "In this tutorial, we've learned how to construct a **textured mesh** from **DensePose model and uv data**, as well as initialize a **Renderer** and change the viewing angle and lighting of our rendered mesh."
   ]
  }
 ],
 "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_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": 2
}