Texturing API updates

Summary:
A fairly big refactor of the texturing API with some breaking changes to how textures are defined.

Main changes:
- There are now 3 types of texture classes: `TexturesUV`, `TexturesAtlas` and `TexturesVertex`. Each class:
   - has a `sample_textures` function which accepts the `fragments` from rasterization and returns `texels`. This means that the shaders will not need to know the type of the mesh texture which will resolve several issues people were reporting on GitHub.
  -  has a `join_batch` method for joining multiple textures of the same type into a batch

Reviewed By: gkioxari

Differential Revision: D21067427

fbshipit-source-id: 4b346500a60181e72fdd1b0dd89b5505c7a33926

Author: nikhilaravi

docs/tutorials/camera_position_optimization_with_differentiable_rendering.ipynb

@@ -520,6 +520,9 @@
  ],
  "metadata": {
   "accelerator": "GPU",
+  "anp_metadata": {
+   "path": "fbsource/fbcode/vision/fair/pytorch3d/docs/tutorials/camera_position_optimization_with_differentiable_rendering.ipynb"
+  },
   "bento_stylesheets": {
    "bento/extensions/flow/main.css": true,
    "bento/extensions/kernel_selector/main.css": true,
@@ -533,6 +536,9 @@
    "provenance": [],
    "toc_visible": true
   },
+  "disseminate_notebook_info": {
+   "backup_notebook_id": "1062179640844868"
+  },
   "kernelspec": {
    "display_name": "pytorch3d (local)",
    "language": "python",

docs/tutorials/render_textured_meshes.ipynb

@@ -84,7 +84,7 @@
     "from skimage.io import imread\n",
     "\n",
     "# Util function for loading meshes\n",
-    "from pytorch3d.io import load_objs_as_meshes\n",
+    "from pytorch3d.io import load_objs_as_meshes, load_obj\n",
     "\n",
     "# Data structures and functions for rendering\n",
     "from pytorch3d.structures import Meshes, Textures\n",
@@ -97,7 +97,7 @@
     "    RasterizationSettings, \n",
     "    MeshRenderer, \n",
     "    MeshRasterizer,  \n",
-    "    TexturedSoftPhongShader\n",
+    "    SoftPhongShader\n",
     ")\n",
     "\n",
     "# add path for demo utils functions \n",
@@ -316,7 +316,7 @@
     "        cameras=cameras, \n",
     "        raster_settings=raster_settings\n",
     "    ),\n",
-    "    shader=TexturedSoftPhongShader(\n",
+    "    shader=SoftPhongShader(\n",
     "        device=device, \n",
     "        cameras=cameras,\n",
     "        lights=lights\n",
@@ -563,6 +563,9 @@
  ],
  "metadata": {
   "accelerator": "GPU",
+  "anp_metadata": {
+   "path": "fbsource/fbcode/vision/fair/pytorch3d/docs/tutorials/render_textured_meshes.ipynb"
+  },
   "bento_stylesheets": {
    "bento/extensions/flow/main.css": true,
    "bento/extensions/kernel_selector/main.css": true,
@@ -575,6 +578,9 @@
    "name": "render_textured_meshes.ipynb",
    "provenance": []
   },
+  "disseminate_notebook_info": {
+   "backup_notebook_id": "569222367081034"
+  },
   "kernelspec": {
    "display_name": "pytorch3d (local)",
    "language": "python",

pytorch3d/datasets/shapenet_base.py

@@ -13,8 +13,8 @@
     OpenGLPerspectiveCameras,
     PointLights,
     RasterizationSettings,
+    TexturesVertex,
 )
-from pytorch3d.structures import Textures
 
 
 class ShapeNetBase(torch.utils.data.Dataset):
@@ -113,8 +113,8 @@ def render(
         """
         paths = self._handle_render_inputs(model_ids, categories, sample_nums, idxs)
         meshes = load_objs_as_meshes(paths, device=device, load_textures=False)
-        meshes.textures = Textures(
-            verts_rgb=torch.ones_like(meshes.verts_padded(), device=device)
+        meshes.textures = TexturesVertex(
+            verts_features=torch.ones_like(meshes.verts_padded(), device=device)
         )
         cameras = kwargs.get("cameras", OpenGLPerspectiveCameras()).to(device)
         renderer = MeshRenderer(

pytorch3d/io/obj_io.py

@@ -11,7 +11,8 @@
 import torch
 from pytorch3d.io.mtl_io import load_mtl, make_mesh_texture_atlas
 from pytorch3d.io.utils import _open_file
-from pytorch3d.structures import Meshes, Textures, join_meshes_as_batch
+from pytorch3d.renderer import TexturesAtlas, TexturesUV
+from pytorch3d.structures import Meshes, join_meshes_as_batch
 
 
 def _make_tensor(data, cols: int, dtype: torch.dtype, device="cpu") -> torch.Tensor:
@@ -41,6 +42,10 @@ def _format_faces_indices(faces_indices, max_index, device, pad_value=None):
     Args:
         faces_indices: List of ints of indices.
         max_index: Max index for the face property.
+        pad_value: if any of the face_indices are padded, specify
+            the value of the padding (e.g. -1). This is only used
+            for texture indices indices where there might
+            not be texture information for all the faces.
 
     Returns:
         faces_indices: List of ints of indices.
@@ -65,7 +70,9 @@ def _format_faces_indices(faces_indices, max_index, device, pad_value=None):
         faces_indices[mask] = pad_value
 
     # Check indices are valid.
-    if torch.any(faces_indices >= max_index) or torch.any(faces_indices < 0):
+    if torch.any(faces_indices >= max_index) or (
+        pad_value is None and torch.any(faces_indices < 0)
+    ):
         warnings.warn("Faces have invalid indices")
 
     return faces_indices
@@ -227,7 +234,14 @@ def load_obj(
         )
 
 
-def load_objs_as_meshes(files: list, device=None, load_textures: bool = True):
+def load_objs_as_meshes(
+    files: list,
+    device=None,
+    load_textures: bool = True,
+    create_texture_atlas: bool = False,
+    texture_atlas_size: int = 4,
+    texture_wrap: Optional[str] = "repeat",
+):
     """
     Load meshes from a list of .obj files using the load_obj function, and
     return them as a Meshes object. This only works for meshes which have a
@@ -246,18 +260,31 @@ def load_objs_as_meshes(files: list, device=None, load_textures: bool = True):
     """
     mesh_list = []
     for f_obj in files:
-        # TODO: update this function to support the two texturing options.
-        verts, faces, aux = load_obj(f_obj, load_textures=load_textures)
-        verts = verts.to(device)
+        verts, faces, aux = load_obj(
+            f_obj,
+            load_textures=load_textures,
+            create_texture_atlas=create_texture_atlas,
+            texture_atlas_size=texture_atlas_size,
+            texture_wrap=texture_wrap,
+        )
         tex = None
-        tex_maps = aux.texture_images
-        if tex_maps is not None and len(tex_maps) > 0:
-            verts_uvs = aux.verts_uvs[None, ...].to(device)  # (1, V, 2)
-            faces_uvs = faces.textures_idx[None, ...].to(device)  # (1, F, 3)
-            image = list(tex_maps.values())[0].to(device)[None]
-            tex = Textures(verts_uvs=verts_uvs, faces_uvs=faces_uvs, maps=image)
-
-        mesh = Meshes(verts=[verts], faces=[faces.verts_idx.to(device)], textures=tex)
+        if create_texture_atlas:
+            # TexturesAtlas type
+            tex = TexturesAtlas(atlas=[aux.texture_atlas])
+        else:
+            # TexturesUV type
+            tex_maps = aux.texture_images
+            if tex_maps is not None and len(tex_maps) > 0:
+                verts_uvs = aux.verts_uvs.to(device)  # (V, 2)
+                faces_uvs = faces.textures_idx.to(device)  # (F, 3)
+                image = list(tex_maps.values())[0].to(device)[None]
+                tex = TexturesUV(
+                    verts_uvs=[verts_uvs], faces_uvs=[faces_uvs], maps=image
+                )
+
+        mesh = Meshes(
+            verts=[verts.to(device)], faces=[faces.verts_idx.to(device)], textures=tex
+        )
         mesh_list.append(mesh)
     if len(mesh_list) == 1:
         return mesh_list[0]

pytorch3d/renderer/__init__.py

@@ -28,11 +28,11 @@
     SoftGouraudShader,
     SoftPhongShader,
     SoftSilhouetteShader,
-    TexturedSoftPhongShader,
+    Textures,
+    TexturesAtlas,
+    TexturesUV,
+    TexturesVertex,
     gouraud_shading,
-    interpolate_face_attributes,
-    interpolate_texture_map,
-    interpolate_vertex_colors,
     phong_shading,
     rasterize_meshes,
 )

pytorch3d/renderer/mesh/__init__.py

@@ -1,21 +1,21 @@
 # Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
 
 
-from .texturing import interpolate_texture_map, interpolate_vertex_colors  # isort:skip
 from .rasterize_meshes import rasterize_meshes
 from .rasterizer import MeshRasterizer, RasterizationSettings
 from .renderer import MeshRenderer
+from .shader import TexturedSoftPhongShader  # DEPRECATED
 from .shader import (
     HardFlatShader,
     HardGouraudShader,
     HardPhongShader,
     SoftGouraudShader,
     SoftPhongShader,
     SoftSilhouetteShader,
-    TexturedSoftPhongShader,
 )
 from .shading import gouraud_shading, phong_shading
-from .utils import interpolate_face_attributes
+from .textures import Textures  # DEPRECATED
+from .textures import TexturesAtlas, TexturesUV, TexturesVertex
 
 
 __all__ = [k for k in globals().keys() if not k.startswith("_")]

pytorch3d/renderer/mesh/shader.py

@@ -1,5 +1,6 @@
 # Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
 
+import warnings
 
 import torch
 import torch.nn as nn
@@ -13,7 +14,6 @@
 from ..lighting import PointLights
 from ..materials import Materials
 from .shading import flat_shading, gouraud_shading, phong_shading
-from .texturing import interpolate_texture_map, interpolate_vertex_colors
 
 
 # A Shader should take as input fragments from the output of rasterization
@@ -57,7 +57,7 @@ def forward(self, fragments, meshes, **kwargs) -> torch.Tensor:
                 or in the forward pass of HardPhongShader"
             raise ValueError(msg)
 
-        texels = interpolate_vertex_colors(fragments, meshes)
+        texels = meshes.sample_textures(fragments)
         lights = kwargs.get("lights", self.lights)
         materials = kwargs.get("materials", self.materials)
         blend_params = kwargs.get("blend_params", self.blend_params)
@@ -104,9 +104,11 @@ def forward(self, fragments, meshes, **kwargs) -> torch.Tensor:
             msg = "Cameras must be specified either at initialization \
                 or in the forward pass of SoftPhongShader"
             raise ValueError(msg)
-        texels = interpolate_vertex_colors(fragments, meshes)
+
+        texels = meshes.sample_textures(fragments)
         lights = kwargs.get("lights", self.lights)
         materials = kwargs.get("materials", self.materials)
+        blend_params = kwargs.get("blend_params", self.blend_params)
         colors = phong_shading(
             meshes=meshes,
             fragments=fragments,
@@ -115,7 +117,7 @@ def forward(self, fragments, meshes, **kwargs) -> torch.Tensor:
             cameras=cameras,
             materials=materials,
         )
-        images = softmax_rgb_blend(colors, fragments, self.blend_params)
+        images = softmax_rgb_blend(colors, fragments, blend_params)
         return images
 
 
@@ -154,6 +156,12 @@ def forward(self, fragments, meshes, **kwargs) -> torch.Tensor:
         lights = kwargs.get("lights", self.lights)
         materials = kwargs.get("materials", self.materials)
         blend_params = kwargs.get("blend_params", self.blend_params)
+
+        # As Gouraud shading applies the illumination to the vertex
+        # colors, the interpolated pixel texture is calculated in the
+        # shading step. In comparison, for Phong shading, the pixel
+        # textures are computed first after which the illumination is
+        # applied.
         pixel_colors = gouraud_shading(
             meshes=meshes,
             fragments=fragments,
@@ -210,54 +218,25 @@ def forward(self, fragments, meshes, **kwargs) -> torch.Tensor:
         return images
 
 
-class TexturedSoftPhongShader(nn.Module):
+def TexturedSoftPhongShader(
+    device="cpu", cameras=None, lights=None, materials=None, blend_params=None
+):
     """
-    Per pixel lighting applied to a texture map. First interpolate the vertex
-    uv coordinates and sample from a texture map. Then apply the lighting model
-    using the interpolated coords and normals for each pixel.
-
-    The blending function returns the soft aggregated color using all
-    the faces per pixel.
-
-    To use the default values, simply initialize the shader with the desired
-    device e.g.
-
-    .. code-block::
-
-        shader = TexturedPhongShader(device=torch.device("cuda:0"))
+    TexturedSoftPhongShader class has been DEPRECATED. Use SoftPhongShader instead.
+    Preserving TexturedSoftPhongShader as a function for backwards compatibility.
     """
-
-    def __init__(
-        self, device="cpu", cameras=None, lights=None, materials=None, blend_params=None
-    ):
-        super().__init__()
-        self.lights = lights if lights is not None else PointLights(device=device)
-        self.materials = (
-            materials if materials is not None else Materials(device=device)
-        )
-        self.cameras = cameras
-        self.blend_params = blend_params if blend_params is not None else BlendParams()
-
-    def forward(self, fragments, meshes, **kwargs) -> torch.Tensor:
-        cameras = kwargs.get("cameras", self.cameras)
-        if cameras is None:
-            msg = "Cameras must be specified either at initialization \
-                or in the forward pass of TexturedSoftPhongShader"
-            raise ValueError(msg)
-        texels = interpolate_texture_map(fragments, meshes)
-        lights = kwargs.get("lights", self.lights)
-        materials = kwargs.get("materials", self.materials)
-        blend_params = kwargs.get("blend_params", self.blend_params)
-        colors = phong_shading(
-            meshes=meshes,
-            fragments=fragments,
-            texels=texels,
-            lights=lights,
-            cameras=cameras,
-            materials=materials,
-        )
-        images = softmax_rgb_blend(colors, fragments, blend_params)
-        return images
+    warnings.warn(
+        """TexturedSoftPhongShader is now deprecated;
+            use SoftPhongShader instead.""",
+        PendingDeprecationWarning,
+    )
+    return SoftPhongShader(
+        device=device,
+        cameras=cameras,
+        lights=lights,
+        materials=materials,
+        blend_params=blend_params,
+    )
 
 
 class HardFlatShader(nn.Module):
@@ -291,7 +270,7 @@ def forward(self, fragments, meshes, **kwargs) -> torch.Tensor:
             msg = "Cameras must be specified either at initialization \
                 or in the forward pass of HardFlatShader"
             raise ValueError(msg)
-        texels = interpolate_vertex_colors(fragments, meshes)
+        texels = meshes.sample_textures(fragments)
         lights = kwargs.get("lights", self.lights)
         materials = kwargs.get("materials", self.materials)
         blend_params = kwargs.get("blend_params", self.blend_params)