temp (#563)

Signed-off-by: JerryJiehanWang <jiehanw@nvidia.com> initial commit Signed-off-by: JerryJiehanWang <jiehanw@nvidia.com> fix copyright Signed-off-by: JerryJiehanWang <jiehanw@nvidia.com> add golden file Signed-off-by: JerryJiehanWang <jiehanw@nvidia.com> temp Signed-off-by: JerryJiehanWang <jiehanw@nvidia.com> uniforming some common functions Signed-off-by: JerryJiehanWang <jiehanw@nvidia.com> add io.utils.py Signed-off-by: JerryJiehanWang <jiehanw@nvidia.com> fix according to review Signed-off-by: JerryJiehanWang <jiehanw@nvidia.com> fix according to review Signed-off-by: JerryJiehanWang <jiehanw@nvidia.com> fix Signed-off-by: JerryJiehanWang <jiehanw@nvidia.com> fix docs Signed-off-by: JerryJiehanWang <jiehanw@nvidia.com> fix docs Signed-off-by: JerryJiehanWang <jiehanw@nvidia.com>
NVIDIAGameWorks · May 12, 2022 · 861e7d7 · 861e7d7
1 parent 3a44a84
commit 861e7d7
Show file tree

Hide file tree

Showing 8 changed files with 252 additions and 36 deletions.
diff --git a/docs/modules/kaolin.io.rst b/docs/modules/kaolin.io.rst
@@ -22,3 +22,4 @@ and :ref:`materials module<kaolin.io.materials>` contains Materials definition t
    kaolin.io.usd
    kaolin.io.modelnet
    kaolin.io.shrec
+   kaolin.io.utils
diff --git a/kaolin/io/__init__.py b/kaolin/io/__init__.py
@@ -7,3 +7,4 @@
 from . import usd
 from . import modelnet
 from . import shrec
+from . import utils
diff --git a/kaolin/io/obj.py b/kaolin/io/obj.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2019,20-21 NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2019,20-21-22 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -21,6 +21,7 @@
 from PIL import Image
 
 from kaolin.io.materials import MaterialLoadError, MaterialFileError, MaterialNotFoundError
+from kaolin.io import utils
 
 __all__ = [
     'ignore_error_handler',
@@ -33,7 +34,6 @@
                          ['vertices', 'faces', 'uvs', 'face_uvs_idx', 'materials',
                           'materials_order', 'vertex_normals', 'face_normals'])
 
-
 def ignore_error_handler(error, **kwargs):
     """Simple error handler to use in :func:`load_obj` that ignore all errors"""
     pass
@@ -49,10 +49,18 @@ def default_error_handler(error, **kwargs):
     """Simple error handle to use in :func:`load_obj` that raises all errors."""
     raise error
 
+def flatten_feature(feature):
+    """Flatten the nested list of a feature.
+    """
+    if feature is None or len(feature) == 0:
+        return None
+    else:
+        return [item for sublist in feature for item in sublist]
+
 
 # TODO(cfujitsang): support https://en.wikipedia.org/wiki/Wavefront_.obj_file#Geometric_vertex ?
 def import_mesh(path, with_materials=False, with_normals=False,
-                error_handler=None):
+                error_handler=None, heterogeneous_mesh_handler=None):
     r"""Load data from an obj file as a single mesh.
 
     With limited materials support to Kd, Ka, Ks, map_Kd, map_Ka and map_Ks.
@@ -68,6 +76,11 @@ def import_mesh(path, with_materials=False, with_normals=False,
             ``error_handler(error: Exception, material_name: Optional[str],
             materials: Optional[list[dict]], materials_order: Optional[list])``.
             Default: raise all errors.
+        heterogeneous_mesh_handler (Callable, optional):
+            function that handles the import of heterogeneous mesh,
+            with following signature:
+            ``heterogeneous_mesh_handler(vertices, face_vertex_counts, *args)``
+        Default: Heterogenenous mesh will raise a NonHomogeneousError.
 
     Returns:
         (obj.return_type):
@@ -131,12 +144,12 @@ def import_mesh(path, with_materials=False, with_normals=False,
                     if len(data[1]) > 1 and data[1][1] != '':
                         face_uvs_idx.append([int(d[1]) for d in data])
                     else:
-                        face_uvs_idx.append([0, 0, 0])
+                        face_uvs_idx.append([0] * len(data))
                 if with_normals:
                     if len(data[1]) > 2:
                         face_normals.append([int(d[2]) for d in data])
                     else:
-                        face_normals.append([0, 0, 0])
+                        face_normals.append([0] * len(data))
             elif with_materials and data[0] == 'usemtl':
                 material_name = data[1]
                 if material_name not in materials_idx:
@@ -159,6 +172,24 @@ def import_mesh(path, with_materials=False, with_normals=False,
 
     vertices = torch.FloatTensor([float(el) for sublist in vertices
                                   for el in sublist]).view(-1, 3)
+
+    face_vertex_counts = torch.IntTensor([len(f) for f in faces])
+
+    if not torch.all(face_vertex_counts == face_vertex_counts[0]):
+        if heterogeneous_mesh_handler is None:
+            raise utils.NonHomogeneousMeshError(f'Mesh is non-homogeneous '
+                                                f'and cannot be imported from {path}.'
+                                                f'User can set heterogeneous_mesh_handler.'
+                                                f'See kaolin.io.utils for the available options')
+        else:
+            all_features = [faces, face_uvs_idx, face_normals]
+            # Flatten all features
+            all_features = [flatten_feature(f) for f in all_features]
+
+            mesh = heterogeneous_mesh_handler(vertices, face_vertex_counts, *all_features)
+        if mesh is not None:
+            vertices, face_vertex_counts, faces, face_uvs_idx, face_normals = mesh
+
     faces = torch.LongTensor(faces) - 1
 
     if with_materials:

diff --git a/kaolin/io/usd.py b/kaolin/io/usd.py
@@ -32,27 +32,31 @@
     warnings.warn('Warning: module pxr not found', ImportWarning)
 
 from kaolin.io import materials as usd_materials
-
+from kaolin.io import utils
 
 mesh_return_type = namedtuple('mesh_return_type', ['vertices', 'faces',
                                                    'uvs', 'face_uvs_idx', 'face_normals', 'materials_order',
                                                    'materials'])
 pointcloud_return_type = namedtuple('pointcloud_return_type', ['points', 'colors', 'normals'])
 
-
 class NonHomogeneousMeshError(Exception):
     """Raised when expecting a homogeneous mesh but a heterogenous
     mesh is encountered.
 
+    .. deprecated:: 1.11.0
+       This function is deprecated, use :func:`kaolin.io.utils.NonHomogeneousMeshError`.
+
     Attributes:
         message (str)
     """
     __slots__ = ['message']
 
     def __init__(self, message):
+        warnings.warn("NonHomogeneousMeshError is deprecated, "
+                      "please use kaolin.io.utils.NonHomogeneousMeshError instead", 
+                      DeprecationWarning, stacklevel=2)
         self.message = message
 
-
 def _get_stage_next_free_path(stage, scene_path):
     """Get next free path in the stage."""
     default_prim = stage.GetDefaultPrim()
@@ -381,12 +385,29 @@ def create_stage(file_path, up_axis='Y'):
 # Mesh Functions
 # TODO(jlafleche) Support mesh subgroups for materials
 def heterogeneous_mesh_handler_skip(*args):
-    r"""Skip heterogeneous meshes."""
-    return None
+    r"""Skip heterogeneous meshes.
+
+    .. deprecated:: 1.11.0
+       This function is deprecated, use :func:`kaolin.io.utils.heterogeneous_mesh_handler_skip`.
+
+    """
+    warnings.warn("heterogeneous_mesh_handler_skip is deprecated, "
+                  "please use kaolin.io.utils.heterogeneous_mesh_handler_skip instead", 
+                  DeprecationWarning, stacklevel=2)
 
+    return None
 
 def heterogeneous_mesh_handler_empty(*args):
-    """Return empty tensors for vertices and faces of heterogeneous meshes."""
+    """Return empty tensors for vertices and faces of heterogeneous meshes.
+
+    .. deprecated:: 1.11.0
+       This function is deprecated, use :func:`kaolin.io.utils.heterogeneous_mesh_handler_empty`.
+
+    """
+    warnings.warn("heterogeneous_mesh_handler_empty is deprecated, "
+                  "please use kaolin.io.utils.heterogeneous_mesh_handler_empty instead", 
+                  DeprecationWarning, stacklevel=2)
+
     return (torch.FloatTensor(size=(0, 3)), torch.LongTensor(size=(0,)),
             torch.LongTensor(size=(0, 3)), torch.FloatTensor(size=(0, 2)),
             torch.LongTensor(size=(0, 3)), torch.FloatTensor(size=(0, 3, 3)),
@@ -397,16 +418,23 @@ def heterogeneous_mesh_handler_naive_homogenize(vertices, face_vertex_counts, *a
     r"""Homogenize list of faces containing polygons of varying number of edges to triangles using fan
     triangulation.
 
+    .. deprecated:: 1.11.0
+       This function is deprecated, use :func:`kaolin.io.utils.heterogeneous_mesh_handler_naive_homogenize`.
+
     Args:
         vertices (torch.FloatTensor): Vertices with shape ``(N, 3)``.
         face_vertex_counts (torch.LongTensor): Number of vertices for each face with shape ``(M)``
             for ``M`` faces.
-        faces (torch.LongTensor): Vertex indices for each face of with shape ``(F)`` where ``F`` is
-            the sum of ``face_vertex_counts``.
+        *args: Variable length features that need to be handled. For example, faces and uvs.
 
     Returns:
-        (list of list of int): Homogeneous list of attributes.
+        (list of torch.tensor): Homogeneous list of attributes.
     """
+
+    warnings.warn("heterogeneous_mesh_handler_naive_homogenize is deprecated, "
+                  "please use kaolin.io.utils.heterogeneous_mesh_handler_naive_homogenize instead", 
+                  DeprecationWarning, stacklevel=2)
+
     def _homogenize(attr, face_vertex_counts):
         if attr is not None:
             attr = attr.tolist()
@@ -552,8 +580,8 @@ def import_meshes(file_path, scene_paths=None, with_materials=False, with_normal
         if faces is not None:
             if not torch.all(face_vertex_counts == face_vertex_counts[0]):
                 if heterogeneous_mesh_handler is None:
-                    raise NonHomogeneousMeshError(f'Mesh at {scene_path} is non-homogeneous '
-                                                  f'and cannot be imported from {file_path}.')
+                    raise utils.NonHomogeneousMeshError(f'Mesh at {scene_path} is non-homogeneous '
+                                                        f'and cannot be imported from {file_path}.')
                 else:
                     mesh = heterogeneous_mesh_handler(vertices, face_vertex_counts, faces, uvs,
                                                       face_uvs_idx, face_normals, materials_face_idx)

diff --git a/kaolin/io/utils.py b/kaolin/io/utils.py
@@ -0,0 +1,75 @@
+# Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# 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.
+
+import torch
+
+__all__ = ["heterogeneous_mesh_handler_skip",
+           "heterogeneous_mesh_handler_empty",
+           'heterogeneous_mesh_handler_naive_homogenize',
+           "NonHomogeneousMeshError"]
+
+class NonHomogeneousMeshError(Exception):
+    """Raised when expecting a homogeneous mesh but a heterogenous
+    mesh is encountered.
+    """
+
+    __slots__ = ['message']
+
+    def __init__(self, message):
+        self.message = message
+
+# Mesh Functions
+# TODO(jlafleche) Support mesh subgroups for materials
+def heterogeneous_mesh_handler_skip(*args):
+    r"""Skip heterogeneous meshes."""
+    return None
+
+def heterogeneous_mesh_handler_empty(*args):
+    """Return empty tensors for vertices and faces of heterogeneous meshes."""
+    return (torch.FloatTensor(size=(0, 3)), torch.LongTensor(size=(0,)),
+            torch.LongTensor(size=(0, 3)), torch.FloatTensor(size=(0, 2)),
+            torch.LongTensor(size=(0, 3)), torch.FloatTensor(size=(0, 3, 3)),
+            torch.LongTensor(size=(0,)))
+
+def heterogeneous_mesh_handler_naive_homogenize(vertices, face_vertex_counts, *features):
+    r"""Homogenize list of faces containing polygons of varying number of edges to triangles using fan
+    triangulation.
+
+    Args:
+        vertices (torch.FloatTensor): Vertices with shape ``(N, 3)``.
+        face_vertex_counts (torch.LongTensor): Number of vertices for each face with shape ``(M)``
+            for ``M`` faces.
+        *features: Variable length features that need to be handled. For example, faces and uvs.
+
+    Returns:
+        (list of torch.tensor): Homogeneous list of attributes.
+    """
+    def _homogenize(attr, face_vertex_counts):
+        if attr is not None:
+            attr = attr if isinstance(attr, list) else attr.tolist()
+            idx = 0
+            new_attr = []
+            for face_vertex_count in face_vertex_counts:
+                attr_face = attr[idx:(idx + face_vertex_count)]
+                idx += face_vertex_count
+                while len(attr_face) >= 3:
+                    new_attr.append(attr_face[:3])
+                    attr_face.pop(1)
+            return torch.tensor(new_attr)
+        else:
+            return None
+
+    new_attrs = [_homogenize(a, face_vertex_counts) for a in features]
+    new_counts = torch.ones(vertices.size(0), dtype=torch.long).fill_(3)
+    return (vertices, new_counts, *new_attrs)