Implemented batched SurfaceMesh class that greatly simplifies working…

… with surface meshes (#740)

* Implemented batched SurfaceMesh class that greatly simplifies working with
normals and other attributes loaded from different mesh representations.

Changes also include:
- non-backward compatible changes to USD and OBJ readers
- extensive tutorial for working with meshes
- much simplified boiler plate code in manu tutorials
- utility to compute vertex normals
- Added utility function to center points and simplified code duplicated across tutorials

Signed-off-by: Maria Masha Shugrina <mshugrina@nvidia.com>

* update python in readthedocs

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

* force scipy version to 1.10.1

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

---------

Signed-off-by: Maria Masha Shugrina <mshugrina@nvidia.com>
Signed-off-by: Clement Fuji Tsang <cfujitsang@nvidia.com>
Co-authored-by: Maria Masha Shugrina <mshugrina@nvidia.com>
Co-authored-by: Clement Fuji Tsang <cfujitsang@nvidia.com>

docs/_templates/layout.html

@@ -4,6 +4,7 @@
   <style>
     :root {
       --nvidia-color: #76B900;
+      --dark-green: #008564;
     }
 
     a, a:visited, a:active {
@@ -27,13 +28,35 @@
         background: #b8d27c;
     }
 
+    html.writer-html4 .rst-content dl:not(.docutils) dl:not(.field-list)>dt, html.writer-html5 .rst-content dl[class]:not(.option-list):not(.field-list):not(.footnote):not(.glossary):not(.simple) dl:not(.field-list):not(.simple)>dt.sig {
+        background-color: #eaefe0;
+        border-left: 3px solid var(--nvidia-color);
+    }
+
+    html.writer-html4 .rst-content dl:not(.docutils)>dt, html.writer-html5 .rst-content dl[class]:not(.option-list):not(.field-list):not(.footnote):not(.glossary):not(.simple)>dt.sig {
+      background: #eaefe0;
+      border-top: 3px solid var(--nvidia-color);
+    }
+
+    html.writer-html4 .rst-content dl:not(.docutils)>dt, html.writer-html5 .rst-content dl[class]:not(.option-list):not(.field-list):not(.footnote):not(.glossary):not(.simple)>dt {
+        color: var(--dark-green);
+    }
+
     .icon, .version, a.icon.icon-home {
         color: white;
     }
 
     table.center-align-center-col td {
-    text-align: center
-}
+        text-align: center
+    }
+
+    .rubric, p.rubric {
+        margin-bottom: 15px;
+        font-weight: 700;
+        font-size: 120%;
+        color: var(--dark-green);
+        border-bottom: 1px solid var(--dark-green);
+    }
 
   </style>
 {% endblock %}

docs/conf.py

@@ -38,6 +38,8 @@
 
 todo_include_todos = True
 
+autodoc_typehints = "description"
+
 intersphinx_mapping = {
     'python': ("https://docs.python.org/3", None),
     'numpy': ('https://numpy.org/doc/stable/', None),

docs/modules/kaolin.ops.rst

@@ -14,6 +14,7 @@ Tensor batching operators are in :ref:`kaolin.ops.batch`, conversions of 3D mode
    kaolin.ops.batch
    kaolin.ops.coords
    kaolin.ops.conversions
+   kaolin.ops.pointcloud
    kaolin.ops.gcn
    kaolin.ops.mesh
    kaolin.ops.random

docs/modules/kaolin.render.camera.rst

@@ -3,6 +3,8 @@
 kaolin.render.camera
 ====================
 
+Kaolin provides extensive camera API. For an overview, see the :ref:`Camera class docs <kaolin.render.camera.Camera>`.
+
 API
 ---
 

docs/modules/kaolin.rep.rst

@@ -0,0 +1,27 @@
+.. _kaolin.rep:
+
+kaolin.rep
+==========
+
+This module includes higher-level Kaolin classes ("representations").
+
+API
+---
+
+Classes
+^^^^^^^
+
+* :ref:`SurfaceMesh <kaolin.rep.SurfaceMesh>`
+* :ref:`Spc <kaolin.rep.Spc>`
+
+Other
+^^^^^^^^^
+
+.. automodule:: kaolin.rep
+   :members:
+   :exclude-members:
+       SurfaceMesh,
+       Spc
+   :undoc-members:
+   :show-inheritance:
+

docs/modules/kaolin.rep.spc.rst

@@ -0,0 +1,14 @@
+:orphan:
+
+.. _kaolin.rep.Spc:
+
+kaolin.rep.Spc
+===========================
+
+API
+---
+
+.. autoclass:: kaolin.rep.Spc
+   :members:
+   :undoc-members:
+   :show-inheritance:

docs/modules/kaolin.rep.surface_mesh.rst

@@ -0,0 +1,146 @@
+:orphan:
+
+.. _kaolin.rep.SurfaceMesh:
+
+SurfaceMesh
+===========================
+
+Tutorial
+--------
+
+For a walk-through of :class:`kaolin.rep.SurfaceMesh` features,
+see `working_with_meshes.ipynb <https://github.com/NVIDIAGameWorks/kaolin/blob/master/examples/tutorial/working_with_meshes.ipynb>`_.
+
+API
+---
+
+* :ref:`Overview <rubric mesh overview>`
+* :ref:`Supported Attributes <rubric mesh attributes>`
+* :ref:`Batching <rubric mesh batching>`
+* :ref:`Attribute Access and Auto-Computability <rubric mesh attribute access>`
+* :ref:`Inspecting and Copying <rubric mesh inspecting>`
+* :ref:`Tensor Operations <rubric mesh tensor ops>`
+
+.. autoclass:: kaolin.rep.SurfaceMesh
+   :members:
+   :undoc-members:
+   :member-order: bysource
+   :exclude-members: Batching, attribute_info_string, set_batching, to_batched, getattr_batched, cat,
+      vertices, face_vertices, normals, face_normals, vertex_normals, uvs, face_uvs, faces, face_normals_idx, face_uvs_idx,
+      material_assignments, materials, cuda, cpu, to, float_tensors_to, detach, get_attributes, has_attribute, has_or_can_compute_attribute,
+      probably_can_compute_attribute, get_attribute, get_or_compute_attribute, check_sanity, to_string, as_dict, describe_attribute,
+      unset_attributes_return_none, allow_auto_compute, batching, convert_attribute_batching
+
+
+   .. _rubric mesh batching:
+
+   .. rubric:: Supported Batching Strategies
+
+   ``SurfaceMesh`` can be instantiated with any of the following batching
+   strategies, and supports conversions between batching strategies. Current
+   batching strategy of a ``mesh`` object can be read from ``mesh.batching`` or
+   by running ``print(mesh)``.
+
+   For example::
+
+       mesh = kaolin.io.obj.load_mesh(path)
+       print(mesh)
+       mesh.to_batched()
+       print(mesh)
+
+   .. autoclass:: kaolin.rep.SurfaceMesh.Batching
+      :members:
+
+   .. automethod:: attribute_info_string
+   .. automethod:: check_sanity
+   .. automethod:: set_batching
+   .. automethod:: to_batched
+   .. automethod:: getattr_batched
+   .. automethod:: cat
+   .. automethod:: convert_attribute_batching
+
+   .. _rubric mesh attribute access:
+
+   .. rubric:: Attribute Access
+
+   By default, ``SurfaceMesh`` will attempt to auto-compute missing attributes
+   on access. These attributes will be cached, unless their ancestors have
+   ``requires_grad == True``. This behavior of the ``mesh`` object can be changed
+   at construction time (``allow_auto_compute=False``) or by setting
+   ``mesh.allow_auto_compute`` later. In addition to this convenience API,
+   explicit methods for attribute access are also supported.
+
+   For example, using **convenience API**::
+
+       # Caching is enabled by default
+       mesh = kaolin.io.obj.load_mesh(path, with_normals=False)
+       print(mesh)
+       print(mesh.has_attribute('face_normals'))  # False
+       fnorm = mesh.face_normals  # Auto-computed
+       print(mesh.has_attribute('face_normals'))  # True (cached)
+
+       # Caching is disabled when gradients need to flow
+       mesh = kaolin.io.obj.load_mesh(path, with_normals=False)
+       mesh.vertices.requires_grad = True   # causes caching to be off
+       print(mesh.has_attribute('face_normals'))  # False
+       fnorm = mesh.face_normals  # Auto-computed
+       print(mesh.has_attribute('face_normals'))  # False (caching disabled)
+
+
+   For example, using **explicit API**::
+
+       mesh = kaolin.io.obj.load_mesh(path, with_normals=False)
+       print(mesh.has_attribute('face_normals'))  # False
+       fnorm = mesh.get_or_compute_attribute('face_normals', should_cache=False)
+       print(mesh.has_attribute('face_normals'))  # False
+
+
+   .. automethod:: get_attributes
+   .. automethod:: has_attribute
+   .. automethod:: has_or_can_compute_attribute
+   .. automethod:: probably_can_compute_attribute
+   .. automethod:: get_attribute
+   .. automethod:: get_or_compute_attribute
+
+   .. _rubric mesh inspecting:
+
+   .. rubric:: Inspecting and Copying Meshes
+
+   To make it easier to work with, ``SurfaceMesh`` supports detailed print
+   statements, as well as ``len()``, ``copy()``, ``deepcopy()`` and can be converted
+   to a dictionary.
+
+   Supported operations::
+
+       import copy
+       mesh_copy = copy.copy(mesh)
+       mesh_copy = copy.deepcopy(mesh)
+       batch_size = len(mesh)
+
+       # Print default attributes
+       print(mesh)
+
+       # Print more detailed attributes
+       print(mesh.to_string(detailed=True, print_stats=True))
+
+       # Print specific attribute
+       print(mesh.describe_attribute('vertices'))
+
+   .. automethod:: to_string
+   .. automethod:: describe_attribute
+   .. automethod:: as_dict
+
+   .. _rubric mesh tensor ops:
+
+   .. rubric:: Tensor Operations
+
+   Convenience operations for device and type conversions of some or all member
+   tensors.
+
+   .. automethod:: cuda
+   .. automethod:: cpu
+   .. automethod:: to
+   .. automethod:: float_tensors_to
+   .. automethod:: detach
+
+   .. rubric:: Other

docs/notes/spc_summary.rst

@@ -274,6 +274,6 @@ Functions useful for working with SPCs are available in the following modules:
 
 * :ref:`kaolin.ops.spc<kaolin.ops.spc>` - general explanation and operations
 * :ref:`kaolin.render.spc<kaolin.render.spc>` - rendering utilities
-* :class:`kaolin.rep.Spc` - high-level wrapper.. _kaolin.ops.spc:
+* :class:`kaolin.rep.Spc` - high-level wrapper
 
 

docs/notes/tutorial_index.rst

@@ -10,6 +10,28 @@ point to master.
 Detailed Tutorials
 ------------------
 
+* `Camera and Rasterization <https://github.com/NVIDIAGameWorks/kaolin/blob/master/examples/tutorial/camera_and_rasterization.ipynb>`_: Rasterize ShapeNet mesh with nvdiffrast and camera:
+    * Load ShapeNet mesh
+    * Preprocess mesh and materials
+    * Create a camera with ``from_args()`` general constructor
+    * Render a mesh with multiple materials with nvdiffrast
+    * Move camera and see the resulting rendering
+* `Optimizing Diffuse Lighting <https://github.com/NVIDIAGameWorks/kaolin/blob/master/examples/tutorial/diffuse_lighting.ipynb>`_: Optimize lighting parameters with spherical gaussians and spherical harmonics:
+    * Load an obj mesh with normals and materials
+    * Rasterize the diffuse and specular albedo
+    * Render and optimize diffuse lighting:
+      * Spherical harmonics
+      * Spherical gaussian with inner product implementation
+      * Spherical gaussian with fitted approximation
+* `Optimize Diffuse and Specular Lighting with Spherical Gaussians <https://github.com/NVIDIAGameWorks/kaolin/blob/master/examples/tutorial/sg_specular_lighting.ipynb>`_:
+    * Load an obj mesh with normals and materials
+    * Generate view rays from camera
+    * Rasterize the diffuse and specular albedo
+    * Render and optimize diffuse and specular lighting with spherical gaussians
+* `Working with Surface Meshes <https://github.com/NVIDIAGameWorks/kaolin/blob/master/examples/tutorial/working_with_meshes.ipynb>`_:
+    * loading and constructing :class:`kaolin.rep.SurfaceMesh` objects
+    * batching of meshes
+    * auto-computing common attributes (like ``face_normals``)
 * `Deep Marching Tetrahedra <https://github.com/NVIDIAGameWorks/kaolin/blob/master/examples/tutorial/dmtet_tutorial.ipynb>`_: reconstructs a tetrahedral mesh from point clouds with `DMTet <https://nv-tlabs.github.io/DMTet/>`_, covering:
     * generating data with Omniverse Kaolin App
     * loading point clouds from a ``.usd`` file
@@ -51,24 +73,7 @@ Detailed Tutorials
     * applying marching tetrahedra
     * using Timelapse API for 3D checkpoints
     * visualizing 3D checkpoints using ``kaolin-dash3d``
-* `Camera and Rasterization <https://github.com/NVIDIAGameWorks/kaolin/blob/master/examples/tutorial/camera_and_rasterization.ipynb>`_: Rasterize ShapeNet mesh with nvdiffrast and camera:
-    * Load ShapeNet mesh
-    * Preprocess mesh and materials
-    * Create a camera with ``from_args()`` general constructor
-    * Render a mesh with multiple materials with nvdiffrast
-    * Move camera and see the resulting rendering
-* `Optimizing Diffuse Lighting <https://github.com/NVIDIAGameWorks/kaolin/blob/master/examples/tutorial/diffuse_lighting.ipynb>`_: Optimize lighting parameters with spherical gaussians and spherical harmonics:
-    * Load an obj mesh with normals and materials
-    * Rasterize the diffuse and specular albedo
-    * Render and optimize diffuse lighting:
-      * Spherical harmonics
-      * Spherical gaussian with inner product implementation
-      * Spherical gaussian with fitted approximation
-* `Optimize Diffuse and Specular Lighting with Spherical Gaussians <https://github.com/NVIDIAGameWorks/kaolin/blob/master/examples/tutorial/sg_specular_lighting.ipynb>`_:
-    * Load an obj mesh with normals and materials
-    * Generate view rays from camera
-    * Rasterize the diffuse and specular albedo
-    * Render and optimize diffuse and specular lighting with spherical gaussians
+
 
 Simple Recipes
 --------------

docs/readthedocs_requirements.txt

@@ -1,3 +1,4 @@
 numpy<1.27.0,>=1.19.5
+scipy==1.10.1
 -f https://download.pytorch.org/whl/lts/1.8/torch_lts.html
 torch==1.8.2+cpu

examples/tutorial/bbox_tutorial.ipynb

@@ -1584,13 +1584,13 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "ani.save(\"animation.gif\", writer=animation.PillowWriter())  # optionally save the animation"
+    "# ani.save(\"animation.gif\", writer=animation.PillowWriter())  # optionally save the animation"
    ]
   }
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": "Python 3",
    "language": "python",
    "name": "python3"
   },
@@ -1604,7 +1604,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.7.11"
+   "version": "3.8.13"
   }
  },
  "nbformat": 4,