Fix some documentation issues (+ a small logging issue for nvdiffrast…

… import) (#565)

Signed-off-by: Clement Fuji Tsang <cfujitsang@nvidia.com>

kaolin/io/materials.py

@@ -246,62 +246,96 @@ class PBRMaterial(Material):
     Supports USD Preview Surface (https://graphics.pixar.com/usd/docs/UsdPreviewSurface-Proposal.html),
     a physically based surface material definition.
 
-    Args:
-        diffuse_color (tuple of floats): RGB values for `Diffuse` parameter (typically referred to as `Albedo`
-            in a metallic workflow) in the range of `(0.0, 0.0, 0.0)` to `(1.0, 1.0, 1.0)`. Default value is grey
-            `(0.5, 0.5, 0.5)`.
-        roughness_value (float): Roughness value of specular lobe in range `0.0` to `1.0`. Default value is `0.5`.
+    Parameters:
+        diffuse_color (tuple of floats):
+            RGB values for `Diffuse` parameter (typically referred to as `Albedo`
+            in a metallic workflow) in the range of `(0.0, 0.0, 0.0)` to `(1.0, 1.0, 1.0)`.
+            Default value is grey `(0.5, 0.5, 0.5)`.
+        roughness_value (float):
+            Roughness value of specular lobe in range `0.0` to `1.0`. Default value is `0.5`.
         metallic_value (float):
             Metallic value, typically set to `0.0` for non-metallic and `1.0` for metallic materials. 
             Ignored if `is_specular_workflow` is `True`. Default value is `0.0`.
-        clearcoat_value (float): Second specular lobe amount. Color is hardcoded to white. Default value is `0.0`.
-        clearcoat_roughness_value (float): Roughness for the clearcoat specular lobe in the range `0.0` to `1.0`.
+        clearcoat_value (float):
+            Second specular lobe amount. Color is hardcoded to white. Default value is `0.0`.
+        clearcoat_roughness_value (float):
+            Roughness for the clearcoat specular lobe in the range `0.0` to `1.0`.
             The default value is `0.01`.
-        opacity_value (float): Opacity, with `1.0` fully opaque and `0.0` as transparent with values within this range
+        opacity_value (float):
+            Opacity, with `1.0` fully opaque and `0.0` as transparent with values within this range
             defining a translucent surface. Default value is `0.0`.
-        opacity_treshold (float): Used to create cutouts based on the `opacity_value`. Surfaces with an opacity
+        opacity_treshold (float):
+            Used to create cutouts based on the `opacity_value`. Surfaces with an opacity
             smaller than the `opacity_threshold` will be fully transparent. Default value is `0.0`.
-        ior_value (float): Index of Refraction used with translucent objects and objects with specular components.
+        ior_value (float):
+            Index of Refraction used with translucent objects and objects with specular components.
             Default value is `1.5`.
-        specular_color (tuple of floats): RGB values for `Specular` lobe. Ignored if `is_specular_workflow` is
-            `False`. Default value is white `(0.0, 0.0, 0.0)`.
-        displacement_value (float): Displacement in the direction of the normal. Default is `0.0`
-        diffuse_texture (torch.FloatTensor): Texture for diffuse parameter, of shape `(3, height, width)`.
-        roughness_texture (torch.FloatTensor): Texture for roughness parameter, of shape `(1, height, width)`.
-        metallic_texture (torch.FloatTensor): Texture for metallic parameter, of shape `(1, height, width)`.
+        specular_color (tuple of floats):
+            RGB values for `Specular` lobe. Ignored if `is_specular_workflow` is `False`.
+            Default value is white `(0.0, 0.0, 0.0)`.
+        displacement_value (float):
+            Displacement in the direction of the normal. Default is `0.0`
+        diffuse_texture (torch.FloatTensor):
+            Texture for diffuse parameter, of shape `(3, height, width)`.
+        roughness_texture (torch.FloatTensor):
+            Texture for roughness parameter, of shape `(1, height, width)`.
+        metallic_texture (torch.FloatTensor):
+            Texture for metallic parameter, of shape `(1, height, width)`.
             Ignored if  `is_specular_workflow` is `True`.
-        clearcoat_texture (torch.FloatTensor): Texture for clearcoat parameter, of shape `(1, height, width)`.
-        clearcoat_roughness_texture (torch.FloatTensor): Texture for clearcoat_roughness parameter, of shape
+        clearcoat_texture (torch.FloatTensor):
+            Texture for clearcoat parameter, of shape `(1, height, width)`.
+        clearcoat_roughness_texture (torch.FloatTensor):
+            Texture for clearcoat_roughness parameter, of shape
             `(1, height, width)`.
-        opacity_texture (torch.FloatTensor): Texture for opacity parameter, of shape `(1, height, width)`.
-        ior_texture (torch.FloatTensor): Texture for opacity parameter, of shape `(1, height, width)`.
-        specular_texture (torch.FloatTensor): Texture for specular parameter, of shape `(3, height, width)`.
+        opacity_texture (torch.FloatTensor):
+            Texture for opacity parameter, of shape `(1, height, width)`.
+        ior_texture (torch.FloatTensor):
+            Texture for opacity parameter, of shape `(1, height, width)`.
+        specular_texture (torch.FloatTensor):
+            Texture for specular parameter, of shape `(3, height, width)`.
             Ignored if `is_specular_workflow` is `False`.
-        normals_texture (torch.FloatTensor): Texture for normal mapping of shape `(3, height, width)`. Normals
-            maps create the illusion of fine three-dimensional detail without increasing the number of polygons.
+        normals_texture (torch.FloatTensor):
+            Texture for normal mapping of shape `(3, height, width)`.
+            Normals maps create the illusion of fine three-dimensional
+            detail without increasing the number of polygons.
             Tensor values must be in the range of `[(-1.0, -1.0, -1.0), (1.0, 1.0, 1.0)]`.
-        displacement_texture (torch.FloatTensor): Texture for displacement in the direction of the normal
-            `(1, height, width)`.
-        is_specular_workflow (bool): Determines whether or not to use a specular workflow. Default
-            is `False` (use a metallic workflow).
-        diffuse_colorspace (string): Colorspace of texture, if provided. Select from [auto, raw, sRGB].
-        roughness_colorspace (string): Colorspace of texture, if provided. Select from [auto, raw].
+        displacement_texture (torch.FloatTensor):
+            Texture for displacement in the direction of the normal `(1, height, width)`.
+        is_specular_workflow (bool):
+            Determines whether or not to use a specular workflow.
+            Default is `False` (use a metallic workflow).
+        diffuse_colorspace (string):
+            Colorspace of texture, if provided. Select from [auto, raw, sRGB].
+            Default is 'auto'.
+        roughness_colorspace (string):
+            Colorspace of texture, if provided. Select from [auto, raw].
             Default is 'auto'.
-        metallic_colorspace (string): Colorspace of texture, if provided. Select from [auto, raw].
+        metallic_colorspace (string):
+            Colorspace of texture, if provided. Select from [auto, raw].
             Default is 'auto'.
-        clearcoat_colorspace (string): Colorspace of texture, if provided. Select from [auto, raw].
+        clearcoat_colorspace (string):
+            Colorspace of texture, if provided. Select from [auto, raw].
             Default is 'auto'.
-        clearcoat_roughness_colorspace (string): Colorspace of texture, if provided. Select from [auto, raw].
+        clearcoat_roughness_colorspace (string):
+            Colorspace of texture, if provided. Select from [auto, raw].
             Default is 'auto'.
-        opacity_colorspace (string): Colorspace of texture, if provided. Select from [auto, raw].
+        opacity_colorspace (string):
+            Colorspace of texture, if provided. Select from [auto, raw].
             Default is 'auto'.
-        ior_colorspace (string): Colorspace of texture, if provided. Select from [auto, raw]. Default is 'auto'.
-        specular_colorspace (string): Colorspace of texture, if provided. Select from [auto, raw, sRGB].
+        ior_colorspace (string):
+            Colorspace of texture, if provided. Select from [auto, raw]. Default is 'auto'.
+        specular_colorspace (string):
+            Colorspace of texture, if provided. Select from [auto, raw, sRGB].
             Default is 'auto'.
-        normals_colorspace (string): Colorspace of texture, if provided. Select from [auto, raw, sRGB].
+        normals_colorspace (string):
+            Colorspace of texture, if provided. Select from [auto, raw, sRGB].
             Default is 'auto'.
-        displacement_colorspace (string): Colorspace of texture, if provided. Select from [auto, raw].
+        displacement_colorspace (string):
+            Colorspace of texture, if provided. Select from [auto, raw].
             Default is 'auto'.
+        shaders (dict):
+            Dictionary mapping a shader name to a reader and writer function.
+            (Currently cannot be set).
     """
     def __init__(
         self,

kaolin/io/usd.py

@@ -485,7 +485,7 @@ def import_mesh(file_path, scene_path=None, with_materials=False, with_normals=F
         - **uvs** (torch.FloatTensor): of shape (num_uvs, 2)
         - **face_uvs_idx** (torch.LongTensor): of shape (num_faces, face_size)
         - **face_normals** (torch.FloatTensor): of shape (num_faces, face_size, 3)
-        - **materials** (list of kaolin.io.materials.Material): Material properties (Not yet implemented)
+        - **materials** (list of kaolin.io.materials.Material): Material properties
 
     Example:
         >>> # Create a stage with some meshes
@@ -539,7 +539,7 @@ def import_meshes(file_path, scene_paths=None, with_materials=False, with_normal
         - **uvs** (list of torch.FloatTensor): of shape (num_uvs, 2)
         - **face_uvs_idx** (list of torch.LongTensor): of shape (num_faces, face_size)
         - **face_normals** (list of torch.FloatTensor): of shape (num_faces, face_size, 3)
-        - **materials** (list of kaolin.io.materials.Material): Material properties (Not yet implemented)
+        - **materials** (list of kaolin.io.materials.Material): Material properties
 
     Example:
         >>> # Create a stage with some meshes
@@ -867,7 +867,7 @@ def import_pointclouds(file_path, scene_paths=None, times=None):
         list of namedtuple of:
             - **points** (list of torch.FloatTensor): of shape (num_points, 3)
             - **colors** (list of torch.FloatTensor): of shape (num_points, 3)
-            - **normals** (list of torch.FloatTensor): of shape (num_points, 2)
+            - **normals** (list of torch.FloatTensor): of shape (num_points, 2) (not yet implemented)
 
     Example:
         >>> points = torch.rand(100, 3)

kaolin/metrics/trianglemesh.py

@@ -72,7 +72,7 @@ def point_to_mesh_distance(pointclouds, face_vertices):
 
     for i in range(batch_size):
         if pointclouds.is_cuda:
-            cur_dist, cur_face_idx, cur_dist_type = UnbatchedTriangleDistanceCuda.apply(
+            cur_dist, cur_face_idx, cur_dist_type = _UnbatchedTriangleDistanceCuda.apply(
                 pointclouds[i], face_vertices[i])
         else:
             cur_dist, cur_face_idx, cur_dist_type = _unbatched_naive_point_to_mesh_distance(
@@ -108,7 +108,7 @@ def _is_not_above(vertex, edge, norm, point):
 def _point_at(vertex, edge, proj):
     return vertex + edge * proj.view(-1, 1)
 
-class UnbatchedTriangleDistanceCuda(torch.autograd.Function):
+class _UnbatchedTriangleDistanceCuda(torch.autograd.Function):
     @staticmethod
     def forward(ctx, points, face_vertices):
         num_points = points.shape[0]

kaolin/ops/conversions/tetmesh.py

@@ -112,9 +112,8 @@ def _unbatched_marching_tetrahedra(vertices, tets, sdf, return_tet_idx):
 
 def marching_tetrahedra(vertices, tets, sdf, return_tet_idx=False):
     r"""Convert discrete signed distance fields encoded on tetrahedral grids to triangle 
-    meshes using marching tetrahedra algorithm as described in "An efficient method of 
-    triangulating equi-valued surfaces by using tetrahedral cells":
-    https://search.ieice.org/bin/summary.php?id=e74-d_1_214. The output surface is differentiable with respect to
+    meshes using marching tetrahedra algorithm as described in `An efficient method of 
+    triangulating equi-valued surfaces by using tetrahedral cells`_. The output surface is differentiable with respect to
     input vertex positions and the SDF values. For more details and example usage in learning, see 
     `Deep Marching Tetrahedra\: a Hybrid Representation for High-Resolution 3D Shape Synthesis`_ NeurIPS 2021.
 
@@ -154,7 +153,10 @@ def marching_tetrahedra(vertices, tets, sdf, return_tet_idx=False):
                 [3, 2, 0]])
         >>> tet_idx_list[0]
         tensor([0, 0])
-        
+
+    .. _An efficient method of triangulating equi-valued surfaces by using tetrahedral cells:
+        https://search.ieice.org/bin/summary.php?id=e74-d_1_214
+
     .. _Deep Marching Tetrahedra\: a Hybrid Representation for High-Resolution 3D Shape Synthesis:
             https://arxiv.org/abs/2111.04276
     """

kaolin/ops/mesh/check_sign.py

@@ -59,15 +59,13 @@ def _unbatched_check_sign_cuda(verts, faces, points):
 
 
 def check_sign(verts, faces, points, hash_resolution=512):
-    r"""Checks if a set of points is contained inside a mesh. 
+    r"""Checks if a set of points is contained inside a watertight triangle mesh.
 
-    Each batch takes in v vertices, f faces of a watertight trimesh, 
-    and p points to check if they are inside the mesh. 
     Shoots a ray from each point to be checked
     and calculates the number of intersections 
     between the ray and triangles in the mesh. 
     Uses the parity of the number of intersections
-    to determine if the point is inside the mesh. 
+    to determine if the point is inside the mesh.
 
     Args:
         verts (torch.Tensor):

kaolin/ops/random.py

@@ -61,9 +61,10 @@ def get_state():
     Mostly used in pair with :func:`set_state`.
 
     See also:
-    - https://pytorch.org/docs/stable/generated/torch.get_rng_state.html#torch.get_rng_state
-    - https://docs.python.org/3/library/random.html#random.getstate
-    - https://numpy.org/doc/stable/reference/random/generated/numpy.random.set_state.html#numpy.random.set_state
+
+    * https://pytorch.org/docs/stable/generated/torch.get_rng_state.html#torch.get_rng_state
+    * https://docs.python.org/3/library/random.html#random.getstate
+    * https://numpy.org/doc/stable/reference/random/generated/numpy.random.set_state.html#numpy.random.set_state
 
     Returns:
        (torch.ByteTensor, tuple, tuple):

kaolin/ops/spc/points.py

@@ -34,8 +34,8 @@ def quantize_points(x, level):
     If a point is out of the range :math:`[-1, 1]` it will be clipped to it.
 
     Args:
-        x (torch.FloatTensor): Floating point coordinates,
-                               must be of last dimension 3.
+        x (torch.Tensor): Floating point coordinates,
+                          must be of last dimension 3.
         level (int): Level of the grid
 
     Returns
@@ -53,10 +53,12 @@ def unbatched_points_to_octree(points, level, sorted=False):
 
     Args:
         points (torch.ShortTensor):
-            The Quantized 3d points. This is not exactly like SPC points hierarchies
-            as this is only the data for a specific level.
+            Quantized 3d points. This is not exactly like SPC points hierarchies
+            as this is only the data for a specific level,
+            of shape :math:`(\text{num_points}, 3)`.
         level (int): Max level of octree, and the level of the points.
         sorted (bool): True if the points are unique and sorted in morton order.
+                       Default=False.
 
     Returns:
         (torch.ByteTensor):
@@ -170,8 +172,7 @@ def coords_to_trilinear(coords, points):
     with ``features`` of shape :math:`(\text{num_points}, 8)`
 
     Args:
-        coords (torch.FloatTensor): Floating point 3D points,
-                                    of shape :math:`(\text{num_points}, 3)`.
+        coords (torch.FloatTensor): 3D points, of shape :math:`(\text{num_points}, 3)`.
         points (torch.ShortTensor): Quantized 3D points (the 0th bit of the voxel x is in),
                                     of shape :math:`(\text{num_points}, 3)`.
 

kaolin/ops/spc/spc.py

@@ -48,13 +48,15 @@ def scan_octrees(octrees, lengths):
     Returns:
         (int, torch.IntTensor, torch.IntTensor):
 
-            - An int containing the depth of the octrees.
-
-            - A tensor containing structural information about the batch of structured point cloud hierarchies,
-              see :ref:`pyramids example <spc_pyramids>`.
-            - A tensor containing the exclusive sum of the bit
+            - max_level, an int containing the depth of the octrees.
+            - :ref:`pyramids<spc_pyramids>`, a tensor containing structural information about
+              the batch of structured point cloud hierarchies,
+              of shape :math:`(\text{batch_size}, 2, \text{max_level + 1})`.
+              See :ref:`the documentation <spc_pyramids>` for more details.
+            - :ref:`exsum<spc_exsum>`, a 1D tensor containing the exclusive sum of the bit
               counts of each byte of the individual octrees within the batched input ``octrees`` tensor,
-              see :ref:`exsum <spc_exsum>`.
+              of size :math:(\text{octree_num_bytes} + \text{batch_size})`.
+              See :ref:`the documentation <spc_exsum>` for more details.
 
     .. note::
 
@@ -79,8 +81,10 @@ def generate_points(octrees, pyramids, exsum):
             counts of individual octrees of shape :math:`(k + \text{batch_size})`
 
     Returns:
-        (torch.Tensor):
-            A tensor containing batched point hierachies derived from a batch of octrees.
+        (torch.ShortTensor):
+            A tensor containing batched point hierachies derived from a batch of octrees,
+            of shape :math:`(\text{num_points_at_all_levels}, 3)`.
+            See :ref:`the documentation<spc_points>` for more details
     """
     return _C.ops.spc.generate_points_cuda(octrees.contiguous(),
                                            pyramids.contiguous(),
@@ -138,9 +142,8 @@ def to_dense(point_hierarchies, pyramids, input, level=-1, **kwargs):
         input (torch.FloatTensor):
             Batched tensor of input feature data,
             of shape :math:`(\text{num_inputs}, \text{feature_dim})`.
-            :math:`\text{num_inputs}`,
-            must correspond to number of points in the
-            batched point hierarchy at `level`.
+            With :math:`\text{num_inputs}` corresponding to a number of points in the
+            batched point hierarchy at ``level``.
 
         level (int):
             The level at which the octree points are converted to feature grids.
@@ -168,8 +171,9 @@ def feature_grids_to_spc(feature_grids, masks=None):
             The sparse 3D feature grids, of shape
             :math:`(\text{batch_size}, \text{feature_dim}, X, Y, Z)`
         masks (optional, torch.BoolTensor):
-            The masks showing where are the features.
-            Default: A feature is determined when not full or zeros.
+            The masks showing where are the features,
+            of shape :math:`(\text{batch_size}, X, Y, Z)`.
+            Default: A feature is determined when not full of zeros.
 
     Returns:
         (torch.ByteTensor, torch.IntTensor, torch.Tensor):
@@ -179,7 +183,7 @@ def feature_grids_to_spc(feature_grids, masks=None):
 
                 - The lengths of each octree, of size :math:`(\text{batch_size})`
 
-                - The coalescent features, of same dtype than `feature_grids`,
+                - The coalescent features, of same dtype than ``feature_grids``,
                   of shape :math:`(\text{num_features}, \text{feature_dim})`.
     """
     batch_size = feature_grids.shape[0]