renderer: add support for rendering high dimensional textures for classification/segmentation use cases (#1248)

Summary:
For 3D segmentation problems it's really useful to be able to train the models from multiple viewpoints using Pytorch3D as the renderer. Currently due to hardcoded assumptions in a few spots the mesh renderer only supports rendering RGB (3 dimensional) data. You can encode the classification information as 3 channel data but if you have more than 3 classes you're out of luck.

This relaxes the assumptions to make rendering semantic classes work with `HardFlatShader` and `AmbientLights` with no diffusion/specular. The other shaders/lights don't make any sense for classification since they mutate the texture values in some way.

This only requires changes in `Materials` and `AmbientLights`. The bulk of the code is the unit test.

Pull Request resolved: https://github.com/facebookresearch/pytorch3d/pull/1248

Test Plan: Added unit test that renders a 5 dimensional texture and compare dimensions 2-5 to a stored picture.

Reviewed By: bottler

Differential Revision: D37764610

Pulled By: d4l3k

fbshipit-source-id: 031895724d9318a6f6bab5b31055bb3f438176a5

pytorch3d/renderer/lighting.py

@@ -292,6 +292,9 @@ class AmbientLights(TensorProperties):
     A light object representing the same color of light everywhere.
     By default, this is white, which effectively means lighting is
     not used in rendering.
+
+    Unlike other lights this supports an arbitrary number of channels, not just 3 for RGB.
+    The ambient_color input determines the number of channels.
     """
 
     def __init__(self, *, ambient_color=None, device: Device = "cpu") -> None:
@@ -304,9 +307,11 @@ class AmbientLights(TensorProperties):
             device: Device (as str or torch.device) on which the tensors should be located
 
         The ambient_color if provided, should be
-            - 3 element tuple/list or list of lists
-            - torch tensor of shape (1, 3)
-            - torch tensor of shape (N, 3)
+            - tuple/list of C-element tuples of floats
+            - torch tensor of shape (1, C)
+            - torch tensor of shape (N, C)
+        where C is the number of channels and N is batch size.
+        For RGB, C is 3.
         """
         if ambient_color is None:
             ambient_color = ((1.0, 1.0, 1.0),)
@@ -317,10 +322,14 @@ class AmbientLights(TensorProperties):
         return super().clone(other)
 
     def diffuse(self, normals, points) -> torch.Tensor:
-        return torch.zeros_like(points)
+        return self._zeros_channels(points)
 
     def specular(self, normals, points, camera_position, shininess) -> torch.Tensor:
-        return torch.zeros_like(points)
+        return self._zeros_channels(points)
+
+    def _zeros_channels(self, points: torch.Tensor) -> torch.Tensor:
+        ch = self.ambient_color.shape[-1]
+        return torch.zeros(*points.shape[:-1], ch, device=points.device)
 
 
 def _validate_light_properties(obj) -> None:

pytorch3d/renderer/materials.py

@@ -27,9 +27,9 @@ class Materials(TensorProperties):
     ) -> None:
         """
         Args:
-            ambient_color: RGB ambient reflectivity of the material
-            diffuse_color: RGB diffuse reflectivity of the material
-            specular_color: RGB specular reflectivity of the material
+            ambient_color: ambient reflectivity of the material
+            diffuse_color: diffuse reflectivity of the material
+            specular_color: specular reflectivity of the material
             shininess: The specular exponent for the material. This defines
                 the focus of the specular highlight with a high value
                 resulting in a concentrated highlight. Shininess values
@@ -37,7 +37,8 @@ class Materials(TensorProperties):
             device: Device (as str or torch.device) on which the tensors should be located
 
         ambient_color, diffuse_color and specular_color can be of shape
-        (1, 3) or (N, 3). shininess can be of shape (1) or (N).
+        (1, C) or (N, C) where C is typically 3 (for RGB). shininess can be of shape (1,)
+        or (N,).
 
         The colors and shininess are broadcast against each other so need to
         have either the same batch dimension or batch dimension = 1.
@@ -49,11 +50,12 @@ class Materials(TensorProperties):
             specular_color=specular_color,
             shininess=shininess,
         )
+        C = self.ambient_color.shape[-1]
         for n in ["ambient_color", "diffuse_color", "specular_color"]:
             t = getattr(self, n)
-            if t.shape[-1] != 3:
-                msg = "Expected %s to have shape (N, 3); got %r"
-                raise ValueError(msg % (n, t.shape))
+            if t.shape[-1] != C:
+                msg = "Expected %s to have shape (N, %d); got %r"
+                raise ValueError(msg % (n, C, t.shape))
         if self.shininess.shape != torch.Size([self._N]):
             msg = "shininess should have shape (N); got %r"
             raise ValueError(msg % repr(self.shininess.shape))