From 6c4151a820658168dd56b85f21b60e6bcf599c92 Mon Sep 17 00:00:00 2001
From: Jeremy Reizenstein <reizenstein@fb.com>
Date: Wed, 17 Mar 2021 16:12:35 -0700
Subject: [PATCH] Remove _read_image from densepose example

Summary:
As noted in #601, the example notebook was using an internal function _read_image from PyTorch3D, which has changed signature recently. It is not meant to be used externally. Switch to using PIL directly.

Other changes: (1) removed unused skimage import. (2) some small tidyups. We now don't have places where cells modify values set by other cells. (3) removed bad calls to `plt.grid` which have no effect.

Reviewed By: theschnitz, nikhilaravi

Differential Revision: D27080372

fbshipit-source-id: 2fce651b3e5d7a4619f0a2b298c5db18c8fa1e2c
---
 docs/tutorials/render_densepose.ipynb | 32 +++++++++++++--------------
 1 file changed, 15 insertions(+), 17 deletions(-)

diff --git a/docs/tutorials/render_densepose.ipynb b/docs/tutorials/render_densepose.ipynb
index 250db70a..fe2fbf1b 100644
--- a/docs/tutorials/render_densepose.ipynb
+++ b/docs/tutorials/render_densepose.ipynb
@@ -92,12 +92,11 @@
     "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",
+    "from PIL import Image\n",
     "import pickle\n",
     "\n",
     "# Data structures and functions for rendering\n",
@@ -184,13 +183,15 @@
     "    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",
+    "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, 1)\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()"
+    "face_indices = torch.Tensor(ALP_UV['All_FaceIndices']).squeeze()  # (13774,)"
    ]
   },
   {
@@ -202,7 +203,6 @@
     "# 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\");"
    ]
   },
@@ -228,6 +228,9 @@
     "        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",
@@ -238,15 +241,15 @@
     "        # 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",
+    "            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[vert_idx] = (1 - V[vert_idx]) / rows + offset_v\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",
-    "U_norm, V_norm = U, 1 - V"
+    "V_norm = 1 - V_norm"
    ]
   },
   {
@@ -263,10 +266,7 @@
     "# 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)"
+    "v_template_extended = v_template[verts-1][None] # (1, 7829, 3)"
    ]
   },
   {
@@ -286,8 +286,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "texture = TexturesUV(maps=tex, faces_uvs=faces, verts_uvs=verts_uv)\n",
-    "mesh = Meshes(v_template_extended, faces, texture)"
+    "texture = TexturesUV(maps=tex, faces_uvs=faces[None], verts_uvs=verts_uv)\n",
+    "mesh = Meshes(v_template_extended, faces[None], texture)"
    ]
   },
   {
@@ -352,7 +352,6 @@
     "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\");"
    ]
   },
@@ -393,7 +392,6 @@
    "source": [
     "plt.figure(figsize=(10, 10))\n",
     "plt.imshow(images[0, ..., :3].cpu().numpy())\n",
-    "plt.grid(\"off\");\n",
     "plt.axis(\"off\");"
    ]
   },