Skip to content

GitLab

"vscode:/vscode.git/clone" did not exist on "c72aae78e6f0d6404eae382defabd25b2f3fe7ae"
Commit f34f4073 authored by Amitav Baruah's avatar Amitav Baruah Committed by Facebook GitHub Bot
Browse files

Add tutorial notebook for rendering densepose 

Summary:
Add a notebook demonstrating how to use Pytorch3D to render a textured mesh with the DensePose textures and the SMPL model

{F336408690}

Reviewed By: nikhilaravi

Differential Revision: D23784314

fbshipit-source-id: c92f32fb9b9468eb7ec26bf58dcabb1f26d92e7b
parent ea0e762c
Show whitespace changes
Inline Side-by-side
Showing with 422 additions and 0 deletions
docs/tutorials/render_densepose.ipynb 0 → 100644
View file @ f34f4073
{
"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 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": [
"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'\n",
"!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'\n",
" \n",
"# 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",
"from skimage.io import imread\n",
"import numpy as np\n",
"\n",
"# libraries for reading data from files\n",
"from scipy.io import loadmat\n",
"from pytorch3d.io.utils import _read_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 http://smpl.is.tue.mpg.de/downloads 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",
"tex = torch.from_numpy(_read_image(file_name=tex_filename, format='RGB') / 255. ).unsqueeze(0).to(device)\n",
"\n",
"verts = torch.from_numpy((ALP_UV[\"All_vertices\"]).astype(int)).squeeze().to(device) # (7829, 1)\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()"
]
},
{
"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.grid(\"off\");\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",
"# 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[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[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",
"U_norm, V_norm = U, 1 - V"
]
},
{
"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 = torch.stack(list(map(lambda vert: v_template[vert-1], verts))).unsqueeze(0).to(device) # (1, 7829, 3)\n",
"\n",
"# add a batch dimension to faces\n",
"faces = faces.unsqueeze(0)"
]
},
{
"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 auxillary 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, verts_uvs=verts_uv)\n",
"mesh = Meshes(v_template_extended, faces, 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.grid(\"off\");\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.grid(\"off\");\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
}
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment