pulsar integration.

Summary:
This diff integrates the pulsar renderer source code into PyTorch3D as an alternative backend for the PyTorch3D point renderer. This diff is the first of a series of three diffs to complete that migration and focuses on the packaging and integration of the source code.

For more information about the pulsar backend, see the release notes and the paper (https://arxiv.org/abs/2004.07484). For information on how to use the backend, see the point cloud rendering notebook and the examples in the folder `docs/examples`.

Tasks addressed in the following diffs:
* Add the PyTorch3D interface,
* Add notebook examples and documentation (or adapt the existing ones to feature both interfaces).

Reviewed By: nikhilaravi

Differential Revision: D23947736

fbshipit-source-id: a5e77b53e6750334db22aefa89b4c079cda1b443

pytorch3d/renderer/points/pulsar/renderer.py

+# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
+"""pulsar renderer PyTorch integration.
+
+Proper Python support for pytorch requires creating a torch.autograd.function
+(independent of whether this is being done within the C++ module). This is done
+here and a torch.nn.Module is exposed for the use in more complex models.
+"""
+import logging
+import math
+import warnings
+from typing import Optional, Tuple, Union
+
+import torch
+
+# pyre-fixme[21]: Could not find a name `_C` defined in module `pytorch3d`.
+from pytorch3d import _C
+from pytorch3d.transforms import axis_angle_to_matrix, rotation_6d_to_matrix
+
+
+LOGGER = logging.getLogger(__name__)
+GAMMA_WARNING_EMITTED = False
+AXANGLE_WARNING_EMITTED = False
+
+
+class _Render(torch.autograd.Function):
+    """
+    Differentiable rendering function for the Pulsar renderer.
+
+    Usually this will be used through the `Renderer` module, which takes care of
+    setting up the buffers and putting them on the correct device. If you use
+    the function directly, you will have to do this manually.
+
+    The steps for this are two-fold: first, you need to create a native Renderer
+    object to provide the required buffers. This is the `native_renderer` parameter
+    for this function. You can create it by creating a `pytorch3d._C.PulsarRenderer`
+    object (with parameters for width, height and maximum number of balls it should
+    be able to render). This object by default resides on the CPU. If you want to
+    shift the buffers to a different device, just assign an empty tensor on the target
+    device to its property `device_tracker`.
+
+    To convert camera parameters from a more convenient representation to the
+    required vectors as in this function, you can use the static
+    function `pytorch3d.renderer.points.pulsar.Renderer._transform_cam_params`.
+
+    Args:
+        * ctx: Pytorch context.
+        * vert_pos: vertex positions. [Bx]Nx3 tensor of positions in 3D space.
+        * vert_col: vertex colors. [Bx]NxK tensor of channels.
+        * vert_rad: vertex radii. [Bx]N tensor of radiuses, >0.
+        * cam_pos: camera position(s). [Bx]3 tensor in 3D coordinates.
+        * pixel_0_0_center: [Bx]3 tensor center(s) of the upper left pixel(s) in
+                            world coordinates.
+        * pixel_vec_x: [Bx]3 tensor from one pixel center to the next in image x
+                       direction in world coordinates.
+        * pixel_vec_y: [Bx]3 tensor from one pixel center to the next in image y
+                       direction in world coordinates.
+        * focal_length: [Bx]1 tensor of focal lengths in world coordinates.
+        * principal_point_offsets: [Bx]2 tensor of principal point offsets in pixels.
+        * gamma: sphere transparency in [1.,1E-5], with 1 being mostly transparent.
+                 [Bx]1.
+        * max_depth: maximum depth for spheres to render. Set this as tighly
+                     as possible to have good numerical accuracy for gradients.
+        * native_renderer: a `pytorch3d._C.PulsarRenderer` object.
+        * min_depth: a float with the minimum depth a sphere must have to be renderer.
+                     Must be 0. or > max(focal_length).
+        * bg_col: K tensor with a background color to use or None (uses all ones).
+        * opacity: [Bx]N tensor of opacity values in [0., 1.] or None (uses all ones).
+        * percent_allowed_difference: a float in [0., 1.[ with the maximum allowed
+                     difference in color space. This is used to speed up the
+                     computation. Default: 0.01.
+        * max_n_hits: a hard limit on the number of hits per ray. Default: max int.
+        * mode: render mode in {0, 1}. 0: render an image; 1: render the hit map.
+        * return_forward_info: whether to return a second map. This second map contains
+            13 channels: first channel contains sm_m (the maximum exponent factor
+            observed), the second sm_d (the normalization denominator, the sum of all
+            coefficients), the third the maximum closest possible intersection for a
+            hit. The following channels alternate with the float encoded integer index
+            of a sphere and its weight. They are the five spheres with the highest
+            color contribution to this pixel color, ordered descending.
+
+        Returns:
+            * image: [Bx]HxWxK float tensor with the resulting image.
+            * forw_info: [Bx]HxWx13 float forward information as described above,
+                  if enabled.
+    """
+
+    @staticmethod
+    def forward(
+        ctx,
+        vert_pos,
+        vert_col,
+        vert_rad,
+        cam_pos,
+        pixel_0_0_center,
+        pixel_vec_x,
+        pixel_vec_y,
+        focal_length,
+        principal_point_offsets,
+        gamma,
+        max_depth,
+        native_renderer,
+        min_depth=0.0,
+        bg_col=None,
+        opacity=None,
+        percent_allowed_difference=0.01,
+        max_n_hits=_C.MAX_UINT,
+        mode=0,
+        return_forward_info=False,
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        if mode != 0:
+            assert not return_forward_info, (
+                "You are using a non-standard rendering mode. This does "
+                "not provide gradients, and also no `forward_info`. Please "
+                "set `return_forward_info` to `False`."
+            )
+        ctx.gamma = gamma
+        ctx.max_depth = max_depth
+        ctx.min_depth = min_depth
+        ctx.percent_allowed_difference = percent_allowed_difference
+        ctx.max_n_hits = max_n_hits
+        ctx.mode = mode
+        ctx.native_renderer = native_renderer
+        image, info = ctx.native_renderer.forward(
+            vert_pos,
+            vert_col,
+            vert_rad,
+            cam_pos,
+            pixel_0_0_center,
+            pixel_vec_x,
+            pixel_vec_y,
+            focal_length,
+            principal_point_offsets,
+            gamma,
+            max_depth,
+            min_depth,
+            bg_col,
+            opacity,
+            percent_allowed_difference,
+            max_n_hits,
+            mode,
+        )
+        if mode != 0:
+            # Backprop not possible!
+            info = None
+        # Prepare for backprop.
+        ctx.save_for_backward(
+            vert_pos,
+            vert_col,
+            vert_rad,
+            cam_pos,
+            pixel_0_0_center,
+            pixel_vec_x,
+            pixel_vec_y,
+            focal_length,
+            principal_point_offsets,
+            bg_col,
+            opacity,
+            image,
+            info,
+        )
+        if return_forward_info:
+            return image, info
+        else:
+            return image
+
+    @staticmethod
+    def backward(ctx, grad_im, *args):
+        global GAMMA_WARNING_EMITTED
+        (
+            vert_pos,
+            vert_col,
+            vert_rad,
+            cam_pos,
+            pixel_0_0_center,
+            pixel_vec_x,
+            pixel_vec_y,
+            focal_length,
+            principal_point_offsets,
+            bg_col,
+            opacity,
+            image,
+            info,
+        ) = ctx.saved_tensors
+        if (
+            (
+                ctx.needs_input_grad[0]
+                or ctx.needs_input_grad[2]
+                or ctx.needs_input_grad[3]
+                or ctx.needs_input_grad[4]
+                or ctx.needs_input_grad[5]
+                or ctx.needs_input_grad[6]
+                or ctx.needs_input_grad[7]
+            )
+            and ctx.gamma < 1e-3
+            and not GAMMA_WARNING_EMITTED
+        ):
+            warnings.warn(
+                "Optimizing for non-color parameters and having a gamma value < 1E-3! "
+                "This is probably not going to produce usable gradients."
+            )
+            GAMMA_WARNING_EMITTED = True
+        if ctx.mode == 0:
+            (
+                grad_pos,
+                grad_col,
+                grad_rad,
+                grad_cam_pos,
+                grad_pixel_0_0_center,
+                grad_pixel_vec_x,
+                grad_pixel_vec_y,
+                grad_opacity,
+            ) = ctx.native_renderer.backward(
+                grad_im,
+                image,
+                info,
+                vert_pos,
+                vert_col,
+                vert_rad,
+                cam_pos,
+                pixel_0_0_center,
+                pixel_vec_x,
+                pixel_vec_y,
+                focal_length,
+                principal_point_offsets,
+                ctx.gamma,
+                ctx.max_depth,
+                ctx.min_depth,
+                bg_col,
+                opacity,
+                ctx.percent_allowed_difference,
+                ctx.max_n_hits,
+                ctx.mode,
+                ctx.needs_input_grad[0],
+                ctx.needs_input_grad[1],
+                ctx.needs_input_grad[2],
+                ctx.needs_input_grad[3]
+                or ctx.needs_input_grad[4]
+                or ctx.needs_input_grad[5]
+                or ctx.needs_input_grad[6],
+                ctx.needs_input_grad[13],
+                None,  # No debug information provided.
+            )
+        else:
+            raise ValueError(
+                "Performing a backward pass for a "
+                "rendering with `mode != 0`! This is not possible."
+            )
+        return (
+            grad_pos,
+            grad_col,
+            grad_rad,
+            grad_cam_pos,
+            grad_pixel_0_0_center,
+            grad_pixel_vec_x,
+            grad_pixel_vec_y,
+            None,  # focal_length
+            None,  # principal_point_offsets
+            None,  # gamma
+            None,  # max_depth
+            None,  # native_renderer
+            None,  # min_depth
+            None,  # bg_col
+            grad_opacity,
+            None,  # percent_allowed_difference
+            None,  # max_n_hits
+            None,  # mode
+            None,  # return_forward_info
+        )
+
+
+class Renderer(torch.nn.Module):
+    """
+    Differentiable rendering module for the Pulsar renderer.
+
+    Set the maximum number of balls to a reasonable value. It is used to determine
+    several buffer sizes. It is no problem to render less balls than this number,
+    but never more.
+
+    When optimizing for sphere positions, sphere radiuses or camera parameters you
+    have to use higher gamma values (closer to one) and larger sphere sizes: spheres
+    can only 'move' to areas that they cover, and only with higher gamma values exists
+    a gradient w.r.t. their color depending on their position.
+
+    Args:
+        * width: result image width in pixels.
+        * height: result image height in pixels.
+        * max_num_balls: the maximum number of balls this renderer will handle.
+        * orthogonal_projection: use an orthogonal instead of perspective projection.
+            Default: False.
+        * right_handed_system: use a right-handed instead of a left-handed coordinate
+            system. This is relevant for compatibility with other drawing or scanning
+            systems. Pulsar by default assumes a left-handed world and camera coordinate
+            system as known from mathematics with x-axis to the right, y axis up and z
+            axis for increasing depth along the optical axis. In the image coordinate
+            system, only the y axis is pointing down, leading still to a left-handed
+            system. If you set this to True, it is assuming a right-handed world and
+            camera coordinate system with x axis to the right, y axis to the top and
+            z axis decreasing along the optical axis. Again, the image coordinate
+            system has a flipped y axis, remaining a right-handed system.
+            Default: False.
+        * background_normalized_depth: the normalized depth the background is placed
+            at.
+            This is on a scale from 0. to 1. between the specified min and max depth
+            (see the forward function). The value 0. is the most furthest depth whereas
+            1. is the closest. Be careful when setting the background too far front - it
+            may hide elements in your scene. Default: EPS.
+        * n_channels: the number of image content channels to use. This is usually three
+            for regular color representations, but can be a higher or lower number.
+            Default: 3.
+        * n_track: the number of spheres to track for gradient calculation per pixel.
+            Only the closest n_track spheres will receive gradients. Default: 5.
+    """
+
+    def __init__(
+        self,
+        width: int,
+        height: int,
+        max_num_balls: int,
+        orthogonal_projection: bool = False,
+        right_handed_system: bool = False,
+        background_normalized_depth: float = _C.EPS,
+        n_channels: int = 3,
+        n_track: int = 5,
+    ):
+        super(Renderer, self).__init__()
+        # pyre-fixme[16]: Module `pytorch3d` has no attribute `_C`.
+        self._renderer = _C.PulsarRenderer(
+            width,
+            height,
+            max_num_balls,
+            orthogonal_projection,
+            right_handed_system,
+            background_normalized_depth,
+            n_channels,
+            n_track,
+        )
+        self.register_buffer("device_tracker", torch.zeros(1))
+
+    @staticmethod
+    def sphere_ids_from_result_info_nograd(result_info: torch.Tensor) -> torch.Tensor:
+        """
+        Get the sphere IDs from a result info tensor.
+        """
+        if result_info.ndim == 3:
+            return Renderer.sphere_ids_from_result_info_nograd(result_info[None, ...])
+        # pyre-fixme[16]: Module `pytorch3d` has no attribute `_C`.
+        return _C.pulsar_sphere_ids_from_result_info_nograd(result_info)
+
+    @staticmethod
+    def depth_map_from_result_info_nograd(result_info: torch.Tensor) -> torch.Tensor:
+        """
+        Get the depth map from a result info tensor.
+
+        This returns a map of the same size as the image with just one channel
+        containing the closest intersection value at that position. Gradients
+        are not available for this tensor, but do note that you can use
+        `sphere_ids_from_result_info_nograd` to get the IDs of the spheres at
+        each position and directly create a loss on their depth if required.
+
+        The depth map contains -1. at positions where no intersection has
+        been detected.
+        """
+        return result_info[..., 4]
+
+    @staticmethod
+    def _transform_cam_params(
+        cam_params: torch.Tensor,
+        width: int,
+        height: int,
+        orthogonal: bool,
+        right_handed: bool,
+    ) -> Tuple[
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+    ]:
+        """
+        Transform 8 component camera parameter vector(s) to the internal camera
+        representation.
+
+        The input vectors consists of:
+            * 3 components for camera position,
+            * 3 components for camera rotation (three rotation angles) or
+              6 components as described in "On the Continuity of Rotation
+              Representations in Neural Networks" (Zhou et al.),
+            * focal length,
+            * the sensor width in world coordinates,
+            * [optional] the principal point offset in x and y.
+
+        The sensor height is inferred by pixel size and sensor width to obtain
+        quadratic pixels.
+
+        Args:
+            * cam_params: [Bx]{8, 10, 11, 13}, input tensors as described above.
+            * width: number of pixels in x direction.
+            * height: number of pixels in y direction.
+            * orthogonal: bool, whether an orthogonal projection is used
+                  (does not use focal length).
+            * right_handed: bool, whether to use a right handed system
+                  (negative z in camera direction).
+
+        Returns:
+            * pos_vec: the position vector in 3D,
+            * pixel_0_0_center: the center of the upper left pixel in world coordinates,
+            * pixel_vec_x: the step to move one pixel on the image x axis
+                   in world coordinates,
+            * pixel_vec_y: the step to move one pixel on the image y axis
+                   in world coordinates,
+            * focal_length: the focal lengths,
+            * principal_point_offsets: the principal point offsets in x, y.
+        """
+        global AXANGLE_WARNING_EMITTED
+        # Set up all direction vectors, i.e., the sensor direction of all axes.
+        assert width > 0
+        assert height > 0
+        batch_processing = True
+        if cam_params.ndimension() == 1:
+            batch_processing = False
+            cam_params = cam_params[None, :]
+        batch_size = cam_params.size(0)
+        continuous_rep = True
+        if cam_params.shape[1] in [8, 10]:
+            if cam_params.requires_grad and not AXANGLE_WARNING_EMITTED:
+                warnings.warn(
+                    "Using an axis angle representation for camera rotations. "
+                    "This has discontinuities and should not be used for optimization. "
+                    "Alternatively, use a six-component representation as described in "
+                    "'On the Continuity of Rotation Representations in Neural Networks'"
+                    " (Zhou et al.). "
+                    "The `pytorch3d.transforms` module provides "
+                    "facilities for using this representation."
+                )
+                AXANGLE_WARNING_EMITTED = True
+            continuous_rep = False
+        else:
+            assert cam_params.shape[1] in [11, 13]
+        pos_vec: torch.Tensor = cam_params[:, :3]
+        principal_point_offsets: torch.Tensor = torch.zeros(
+            (cam_params.shape[0], 2), dtype=torch.int32, device=cam_params.device
+        )
+        if continuous_rep:
+            rot_vec = cam_params[:, 3:9]
+            focal_length: torch.Tensor = cam_params[:, 9:10]
+            sensor_size_x = cam_params[:, 10:11]
+            if cam_params.shape[1] == 13:
+                principal_point_offsets: torch.Tensor = cam_params[:, 11:13].to(
+                    torch.int32
+                )
+        else:
+            rot_vec = cam_params[:, 3:6]
+            focal_length: torch.Tensor = cam_params[:, 6:7]
+            sensor_size_x = cam_params[:, 7:8]
+            if cam_params.shape[1] == 10:
+                principal_point_offsets: torch.Tensor = cam_params[:, 8:10].to(
+                    torch.int32
+                )
+        # Always get quadratic pixels.
+        pixel_size_x = sensor_size_x / float(width)
+        sensor_size_y = height * pixel_size_x
+        LOGGER.debug(
+            "Camera position: %s, rotation: %s. Focal length: %s.",
+            str(pos_vec),
+            str(rot_vec),
+            str(focal_length),
+        )
+        if continuous_rep:
+            rot_mat = rotation_6d_to_matrix(rot_vec)
+        else:
+            rot_mat = axis_angle_to_matrix(rot_vec)
+        sensor_dir_x = torch.matmul(
+            rot_mat,
+            torch.tensor(
+                [1.0, 0.0, 0.0], dtype=torch.float32, device=rot_mat.device
+            ).repeat(batch_size, 1)[:, :, None],
+        )[:, :, 0]
+        sensor_dir_y = torch.matmul(
+            rot_mat,
+            torch.tensor(
+                [0.0, -1.0, 0.0], dtype=torch.float32, device=rot_mat.device
+            ).repeat(batch_size, 1)[:, :, None],
+        )[:, :, 0]
+        sensor_dir_z = torch.matmul(
+            rot_mat,
+            torch.tensor(
+                [0.0, 0.0, 1.0], dtype=torch.float32, device=rot_mat.device
+            ).repeat(batch_size, 1)[:, :, None],
+        )[:, :, 0]
+        if right_handed:
+            sensor_dir_z *= -1
+        LOGGER.debug(
+            "Sensor direction vectors: %s, %s, %s.",
+            str(sensor_dir_x),
+            str(sensor_dir_y),
+            str(sensor_dir_z),
+        )
+        if orthogonal:
+            sensor_center = pos_vec
+        else:
+            sensor_center = pos_vec + focal_length * sensor_dir_z
+        LOGGER.debug("Sensor center: %s.", str(sensor_center))
+        sensor_luc = (  # Sensor left upper corner.
+            sensor_center
+            - sensor_dir_x * (sensor_size_x / 2.0)
+            - sensor_dir_y * (sensor_size_y / 2.0)
+        )
+        LOGGER.debug("Sensor luc: %s.", str(sensor_luc))
+        pixel_size_x = sensor_size_x / float(width)
+        pixel_size_y = sensor_size_y / float(height)
+        LOGGER.debug(
+            "Pixel sizes (x): %s, (y) %s.", str(pixel_size_x), str(pixel_size_y)
+        )
+        pixel_vec_x: torch.Tensor = sensor_dir_x * pixel_size_x
+        pixel_vec_y: torch.Tensor = sensor_dir_y * pixel_size_y
+        pixel_0_0_center = sensor_luc + 0.5 * pixel_vec_x + 0.5 * pixel_vec_y
+        LOGGER.debug(
+            "Pixel 0 centers: %s, vec x: %s, vec y: %s.",
+            str(pixel_0_0_center),
+            str(pixel_vec_x),
+            str(pixel_vec_y),
+        )
+        if not orthogonal:
+            LOGGER.debug(
+                "Camera horizontal fovs: %s deg.",
+                str(
+                    2.0
+                    * torch.atan(0.5 * sensor_size_x / focal_length)
+                    / math.pi
+                    * 180.0
+                ),
+            )
+            LOGGER.debug(
+                "Camera vertical fovs: %s deg.",
+                str(
+                    2.0
+                    * torch.atan(0.5 * sensor_size_y / focal_length)
+                    / math.pi
+                    * 180.0
+                ),
+            )
+        # Reduce dimension.
+        focal_length: torch.Tensor = focal_length[:, 0]
+        if batch_processing:
+            return (
+                pos_vec,
+                pixel_0_0_center,
+                pixel_vec_x,
+                pixel_vec_y,
+                focal_length,
+                principal_point_offsets,
+            )
+        else:
+            return (
+                pos_vec[0],
+                pixel_0_0_center[0],
+                pixel_vec_x[0],
+                pixel_vec_y[0],
+                focal_length[0],
+                principal_point_offsets[0],
+            )
+
+    def forward(
+        self,
+        vert_pos: torch.Tensor,
+        vert_col: torch.Tensor,
+        vert_rad: torch.Tensor,
+        cam_params: torch.Tensor,
+        gamma: float,
+        max_depth: float,
+        min_depth: float = 0.0,
+        bg_col: Optional[torch.Tensor] = None,
+        opacity: Optional[torch.Tensor] = None,
+        percent_allowed_difference: float = 0.01,
+        max_n_hits: int = _C.MAX_UINT,
+        mode: int = 0,
+        return_forward_info: bool = False,
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, Optional[torch.Tensor]]]:
+        """
+        Rendering pass to create an image from the provided spheres and camera
+        parameters.
+
+        Args:
+            * vert_pos: vertex positions. [Bx]Nx3 tensor of positions in 3D space.
+            * vert_col: vertex colors. [Bx]NxK tensor of channels.
+            * vert_rad: vertex radii. [Bx]N tensor of radiuses, >0.
+            * cam_params: camera parameter(s). [Bx]8 tensor, consisting of:
+                - 3 components for camera position,
+                - 3 components for camera rotation (axis angle representation) or
+                  6 components as described in "On the Continuity of Rotation
+                  Representations in Neural Networks" (Zhou et al.),
+                - focal length,
+                - the sensor width in world coordinates,
+                - [optional] an offset for the principal point in x, y (no gradients).
+            * gamma: sphere transparency in [1.,1E-5], with 1 being mostly transparent.
+                [Bx]1.
+            * max_depth: maximum depth for spheres to render. Set this as tightly
+                        as possible to have good numerical accuracy for gradients.
+                        float > min_depth + eps.
+            * min_depth: a float with the minimum depth a sphere must have to be
+                        rendered. Must be 0. or > max(focal_length) + eps.
+            * bg_col: K tensor with a background color to use or None (uses all ones).
+            * opacity: [Bx]N tensor of opacity values in [0., 1.] or None (uses all
+                    ones).
+            * percent_allowed_difference: a float in [0., 1.[ with the maximum allowed
+                        difference in color space. This is used to speed up the
+                        computation. Default: 0.01.
+            * max_n_hits: a hard limit on the number of hits per ray. Default: max int.
+            * mode: render mode in {0, 1}. 0: render an image; 1: render the hit map.
+            * return_forward_info: whether to return a second map. This second map
+                contains 13 channels: first channel contains sm_m (the maximum
+                exponent factor observed), the second sm_d (the normalization
+                denominator, the sum of all coefficients), the third the maximum closest
+                possible intersection for a hit. The following channels alternate with
+                the float encoded integer index of a sphere and its weight. They are the
+                five spheres with the highest color contribution to this pixel color,
+                ordered descending. Default: False.
+
+        Returns:
+            * image: [Bx]HxWx3 float tensor with the resulting image.
+            * forw_info: [Bx]HxWx13 float forward information as described above, if
+                    enabled.
+        """
+        # The device tracker is registered as buffer.
+        # pyre-fixme[16]: `Renderer` has no attribute `device_tracker`.
+        self._renderer.device_tracker = self.device_tracker
+        (
+            pos_vec,
+            pixel_0_0_center,
+            pixel_vec_x,
+            pixel_vec_y,
+            focal_lengths,
+            principal_point_offsets,
+        ) = Renderer._transform_cam_params(
+            cam_params,
+            self._renderer.width,
+            self._renderer.height,
+            self._renderer.orthogonal,
+            self._renderer.right_handed,
+        )
+        if (
+            focal_lengths.min().item() > 0.0
+            and max_depth > 10_000.0 * focal_lengths.min().item()
+        ):
+            warnings.warn(
+                (
+                    "Extreme ratio of `max_depth` vs. focal length detected "
+                    "(%f vs. %f, ratio: %f). This will likely lead to "
+                    "artifacts due to numerical instabilities."
+                )
+                % (
+                    max_depth,
+                    focal_lengths.min().item(),
+                    max_depth / focal_lengths.min().item(),
+                )
+            )
+        # pyre-fixme[16]: `_Render` has no attribute `apply`.
+        ret_res = _Render.apply(
+            vert_pos,
+            vert_col,
+            vert_rad,
+            pos_vec,
+            pixel_0_0_center,
+            pixel_vec_x,
+            pixel_vec_y,
+            # Focal length and sensor size don't need gradients other than through
+            # `pixel_vec_x` and `pixel_vec_y`. The focal length is only used in the
+            # renderer to determine the projection areas of the balls.
+            focal_lengths,
+            # principal_point_offsets does not receive gradients.
+            principal_point_offsets,
+            gamma,
+            max_depth,
+            self._renderer,
+            min_depth,
+            bg_col,
+            opacity,
+            percent_allowed_difference,
+            max_n_hits,
+            mode,
+            (mode == 0) and return_forward_info,
+        )
+        if return_forward_info and mode != 0:
+            return ret_res, None
+        return ret_res
+
+    def extra_repr(self) -> str:
+        """Extra information to print in pytorch graphs."""
+        return "width={}, height={}, max_num_balls={}".format(
+            self._renderer.width, self._renderer.height, self._renderer.max_num_balls
+        )

pytorch3d/transforms/__init__.py

 # Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
 
+from .external.kornia_angle_axis_to_rotation_matrix import (
+    angle_axis_to_rotation_matrix as axis_angle_to_matrix,
+)
 from .rotation_conversions import (
     euler_angles_to_matrix,
     matrix_to_euler_angles,
     matrix_to_quaternion,
+    matrix_to_rotation_6d,
     quaternion_apply,
     quaternion_invert,
     quaternion_multiply,
@@ -12,6 +16,7 @@ from .rotation_conversions import (
     random_quaternions,
     random_rotation,
     random_rotations,
+    rotation_6d_to_matrix,
     standardize_quaternion,
 )
 from .so3 import (

pytorch3d/transforms/external/__init__.py

+# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.

pytorch3d/transforms/external/kornia_angle_axis_to_rotation_matrix.py

+#!/usr/bin/env python3
+"""
+This file contains the great angle axis to rotation matrix conversion
+from kornia (https://github.com/arraiyopensource/kornia). The license
+can be found in kornia_license.txt.
+
+The method is used unchanged; the documentation has been adjusted
+to match our doc format.
+"""
+import torch
+
+
+def angle_axis_to_rotation_matrix(angle_axis):
+    """Convert 3d vector of axis-angle rotation to 4x4 rotation matrix
+
+    Args:
+        angle_axis (Tensor): tensor of 3d vector of axis-angle rotations.
+
+    Returns:
+        Tensor: tensor of 3x3 rotation matrix.
+
+    Shape:
+        - Input: :math:`(N, 3)`
+        - Output: :math:`(N, 3, 3)`
+
+    Example:
+
+    ..code-block::python
+
+        >>> input = torch.rand(1, 3)  # Nx3
+        >>> output = tgm.angle_axis_to_rotation_matrix(input)  # Nx3x3
+        >>> output = tgm.angle_axis_to_rotation_matrix(input)  # Nx3x3
+    """
+
+    def _compute_rotation_matrix(angle_axis, theta2, eps=1e-6):
+        # We want to be careful to only evaluate the square root if the
+        # norm of the angle_axis vector is greater than zero. Otherwise
+        # we get a division by zero.
+        k_one = 1.0
+        theta = torch.sqrt(theta2)
+        wxyz = angle_axis / (theta + eps)
+        wx, wy, wz = torch.chunk(wxyz, 3, dim=1)
+        cos_theta = torch.cos(theta)
+        sin_theta = torch.sin(theta)
+
+        r00 = cos_theta + wx * wx * (k_one - cos_theta)
+        r10 = wz * sin_theta + wx * wy * (k_one - cos_theta)
+        r20 = -wy * sin_theta + wx * wz * (k_one - cos_theta)
+        r01 = wx * wy * (k_one - cos_theta) - wz * sin_theta
+        r11 = cos_theta + wy * wy * (k_one - cos_theta)
+        r21 = wx * sin_theta + wy * wz * (k_one - cos_theta)
+        r02 = wy * sin_theta + wx * wz * (k_one - cos_theta)
+        r12 = -wx * sin_theta + wy * wz * (k_one - cos_theta)
+        r22 = cos_theta + wz * wz * (k_one - cos_theta)
+        rotation_matrix = torch.cat(
+            [r00, r01, r02, r10, r11, r12, r20, r21, r22], dim=1
+        )
+        return rotation_matrix.view(-1, 3, 3)
+
+    def _compute_rotation_matrix_taylor(angle_axis):
+        rx, ry, rz = torch.chunk(angle_axis, 3, dim=1)
+        k_one = torch.ones_like(rx)
+        rotation_matrix = torch.cat(
+            [k_one, -rz, ry, rz, k_one, -rx, -ry, rx, k_one], dim=1
+        )
+        return rotation_matrix.view(-1, 3, 3)
+
+    # stolen from ceres/rotation.h
+
+    _angle_axis = torch.unsqueeze(angle_axis + 1e-6, dim=1)
+    # _angle_axis.register_hook(lambda grad: pdb.set_trace())
+    # _angle_axis = 1e-6
+    theta2 = torch.matmul(_angle_axis, _angle_axis.transpose(1, 2))
+    theta2 = torch.squeeze(theta2, dim=1)
+
+    # compute rotation matrices
+    rotation_matrix_normal = _compute_rotation_matrix(angle_axis, theta2)
+    rotation_matrix_taylor = _compute_rotation_matrix_taylor(angle_axis)
+
+    # create mask to handle both cases
+    eps = 1e-6
+    mask = (theta2 > eps).view(-1, 1, 1).to(theta2.device)
+    mask_pos = (mask).type_as(theta2)
+    mask_neg = (mask == False).type_as(theta2)  # noqa
+
+    # create output pose matrix
+    batch_size = angle_axis.shape[0]
+    rotation_matrix = torch.eye(3).to(angle_axis.device).type_as(angle_axis)
+    rotation_matrix = rotation_matrix.view(1, 3, 3).repeat(batch_size, 1, 1)
+    # fill output matrix with masked values
+    rotation_matrix[..., :3, :3] = (
+        mask_pos * rotation_matrix_normal + mask_neg * rotation_matrix_taylor
+    )
+    return rotation_matrix.to(angle_axis.device).type_as(angle_axis)  # Nx4x4

pytorch3d/transforms/external/kornia_license.txt

+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

tests/bm_barycentric_clipping.py

@@ -110,3 +110,7 @@ def bm_barycentric_clip() -> None:
 
         benchmark(baryclip_cuda, "BARY_CLIP_CUDA", kwargs_list, warmup_iters=1)
         benchmark(baryclip_pytorch, "BARY_CLIP_PYTORCH", kwargs_list, warmup_iters=1)
+
+
+if __name__ == "__main__":
+    bm_barycentric_clip()

tests/bm_blending.py

@@ -42,3 +42,7 @@ def bm_blending() -> None:
         kwargs_list,
         warmup_iters=1,
     )
+
+
+if __name__ == "__main__":
+    bm_blending()

tests/bm_cameras_alignment.py

@@ -22,3 +22,7 @@ def bm_cameras_alignment() -> None:
         kwargs_list,
         warmup_iters=1,
     )
+
+
+if __name__ == "__main__":
+    bm_cameras_alignment()

tests/bm_chamfer.py

@@ -8,6 +8,8 @@ from test_chamfer import TestChamfer
 
 
 def bm_chamfer() -> None:
+    # Currently disabled.
+    return
     devices = ["cpu"]
     if torch.cuda.is_available():
         devices.append("cuda:0")
@@ -53,3 +55,7 @@ def bm_chamfer() -> None:
                 }
             )
         benchmark(TestChamfer.chamfer_with_init, "CHAMFER", kwargs_list, warmup_iters=1)
+
+
+if __name__ == "__main__":
+    bm_chamfer()

tests/bm_cubify.py

@@ -11,3 +11,7 @@ def bm_cubify() -> None:
         {"batch_size": 16, "V": 32},
     ]
     benchmark(TestCubify.cubify_with_init, "CUBIFY", kwargs_list, warmup_iters=1)
+
+
+if __name__ == "__main__":
+    bm_cubify()

tests/bm_face_areas_normals.py

@@ -37,3 +37,7 @@ def bm_face_areas_normals() -> None:
         kwargs_list,
         warmup_iters=1,
     )
+
+
+if __name__ == "__main__":
+    bm_face_areas_normals()

tests/bm_graph_conv.py

@@ -40,3 +40,7 @@ def bm_graph_conv() -> None:
         kwargs_list,
         warmup_iters=1,
     )
+
+
+if __name__ == "__main__":
+    bm_graph_conv()

tests/bm_interpolate_face_attributes.py

@@ -74,3 +74,7 @@ def bm_interpolate_face_attribues() -> None:
         kwargs_list.append({"N": N, "S": S, "K": K, "F": F, "D": D, "impl": impl})
     benchmark(_bm_forward, "FORWARD", kwargs_list, warmup_iters=3)
     benchmark(_bm_forward_backward, "FORWARD+BACKWARD", kwargs_list, warmup_iters=3)
+
+
+if __name__ == "__main__":
+    bm_interpolate_face_attribues()

tests/bm_knn.py

@@ -24,3 +24,7 @@ def bm_knn() -> None:
     benchmark(TestKNN.knn_square, "KNN_SQUARE", kwargs_list, warmup_iters=1)
 
     benchmark(TestKNN.knn_ragged, "KNN_RAGGED", kwargs_list, warmup_iters=1)
+
+
+if __name__ == "__main__":
+    bm_knn()

tests/bm_lighting.py

@@ -45,3 +45,7 @@ def bm_lighting() -> None:
         kwargs_list.append({"N": N, "S": S, "K": K})
     benchmark(_bm_diffuse_cuda_with_init, "DIFFUSE", kwargs_list, warmup_iters=3)
     benchmark(_bm_specular_cuda_with_init, "SPECULAR", kwargs_list, warmup_iters=3)
+
+
+if __name__ == "__main__":
+    bm_lighting()

tests/bm_main.py

@@ -2,8 +2,10 @@
 # Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
 
 import glob
-import importlib
-from os.path import basename, dirname, isfile, join, sys
+import os
+import subprocess
+import sys
+from os.path import dirname, isfile, join
 
 
 if __name__ == "__main__":
@@ -11,20 +13,22 @@ if __name__ == "__main__":
     if len(sys.argv) > 1:
         # Parse from flags.
         # pyre-ignore[16]
-        module_names = [n for n in sys.argv if n.startswith("bm_")]
+        file_names = [
+            join(dirname(__file__), n) for n in sys.argv if n.startswith("bm_")
+        ]
     else:
         # Get all the benchmark files (starting with "bm_").
         bm_files = glob.glob(join(dirname(__file__), "bm_*.py"))
-        module_names = [
-            basename(f)[:-3]
-            for f in bm_files
-            if isfile(f) and not f.endswith("bm_main.py")
-        ]
+        file_names = sorted(
+            f for f in bm_files if isfile(f) and not f.endswith("bm_main.py")
+        )
 
-    for module_name in module_names:
-        module = importlib.import_module(module_name)
-        for attr in dir(module):
-            # Run all the functions with names "bm_*" in the module.
-            if attr.startswith("bm_"):
-                print("Running benchmarks for " + module_name + "/" + attr + "...")
-                getattr(module, attr)()
+    # Forward all important path information to the subprocesses through the
+    # environment.
+    os.environ["PATH"] = sys.path[0] + ":" + os.environ.get("PATH", "")
+    os.environ["LD_LIBRARY_PATH"] = (
+        sys.path[0] + ":" + os.environ.get("LD_LIBRARY_PATH", "")
+    )
+    os.environ["PYTHONPATH"] = ":".join(sys.path)
+    for file_name in file_names:
+        subprocess.check_call([sys.executable, file_name])

tests/bm_mesh_edge_loss.py

@@ -19,3 +19,7 @@ def bm_mesh_edge_loss() -> None:
     benchmark(
         TestMeshEdgeLoss.mesh_edge_loss, "MESH_EDGE_LOSS", kwargs_list, warmup_iters=1
     )
+
+
+if __name__ == "__main__":
+    bm_mesh_edge_loss()

tests/bm_mesh_io.py

@@ -95,3 +95,7 @@ def bm_save_load() -> None:
         kwargs_list,
         warmup_iters=1,
     )
+
+
+if __name__ == "__main__":
+    bm_save_load()

tests/bm_mesh_laplacian_smoothing.py

@@ -30,3 +30,7 @@ def bm_mesh_laplacian_smoothing() -> None:
         kwargs_list,
         warmup_iters=1,
     )
+
+
+if __name__ == "__main__":
+    bm_mesh_laplacian_smoothing()

tests/bm_mesh_normal_consistency.py

@@ -27,3 +27,7 @@ def bm_mesh_normal_consistency() -> None:
         kwargs_list,
         warmup_iters=1,
     )
+
+
+if __name__ == "__main__":
+    bm_mesh_normal_consistency()