Add voxel_range option for mesh to voxel conversion

Signed-off-by: Kosuke Takeuchi <[email protected]>
NVIDIAGameWorks · Jul 10, 2020 · eb09fad · eb09fad
1 parent ec193ff
commit eb09fad
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 5 deletions.
diff --git a/kaolin/conversions/meshconversions.py b/kaolin/conversions/meshconversions.py
@@ -53,7 +53,8 @@ def trianglemesh_to_pointcloud(mesh: kaolin.rep.Mesh, num_points: int):
 
 
 def trianglemesh_to_voxelgrid(mesh: kaolin.rep.Mesh, resolution: int,
-                              normalize: bool = True, vertex_offset: float = 0.):
+                              normalize: bool = True, vertex_offset: float = 0.,
+                              voxel_range: float = 1.):
     r""" Converts mesh to a voxel model of a given resolution
 
     Args:
@@ -63,6 +64,7 @@ def trianglemesh_to_voxelgrid(mesh: kaolin.rep.Mesh, resolution: int,
             unit cube centered at the origin.
         vertex_offset (float): Offset applied to all vertices after
                                normalizing.
+        voxel_range (float): Range of voxelization.
 
     Returns:
         voxels (torch.Tensor): voxel array of desired resolution
@@ -80,6 +82,7 @@ def trianglemesh_to_voxelgrid(mesh: kaolin.rep.Mesh, resolution: int,
         mesh.vertices = (mesh.vertices - verts_min) / (verts_max - verts_min) - 0.5
 
     mesh.vertices = mesh.vertices + vertex_offset
+    mesh.vertices /= voxel_range
 
     points = mesh.vertices
     smallest_side = (1. / resolution)**2

diff --git a/kaolin/datasets/shapenet.py b/kaolin/datasets/shapenet.py
@@ -367,6 +367,7 @@ class ShapeNet_Voxels(data.Dataset):
         split (float): amount of dataset that is training out of 1
         resolutions (list): list of resolutions to be returned
         no_progress (bool): if True, disables progress bar
+        voxel_range (float): Range of voxelization.
 
     Returns:
         .. code-block::
@@ -387,7 +388,8 @@ class ShapeNet_Voxels(data.Dataset):
     """
 
     def __init__(self, root: str, cache_dir: str, categories: list = ['chair'], train: bool = True,
-                 split: float = .7, resolutions=[128, 32], no_progress: bool = False):
+                 split: float = .7, resolutions=[128, 32], no_progress: bool = False,
+                 voxel_range: float = 1.0):
         self.root = Path(root)
         self.cache_dir = Path(cache_dir) / 'voxels'
         self.cache_transforms = {}
@@ -406,7 +408,10 @@ def __init__(self, root: str, cache_dir: str, categories: list = ['chair'], trai
 
         for res in self.params['resolutions']:
             self.cache_transforms[res] = tfs.CacheCompose([
-                tfs.TriangleMeshToVoxelGrid(res, normalize=False, vertex_offset=0.5),
+                tfs.TriangleMeshToVoxelGrid(res,
+                                            normalize=False,
+                                            vertex_offset=0.5,
+                                            voxel_range=voxel_range),
                 tfs.FillVoxelGrid(thresh=0.5),
                 tfs.ExtractProjectOdmsFromVoxelGrid()
             ], self.cache_dir)

diff --git a/kaolin/transforms/transforms.py b/kaolin/transforms/transforms.py
@@ -662,15 +662,18 @@ class TriangleMeshToVoxelGrid(Transform):
             unit cube centered at the origin.
         vertex_offset (float): Offset applied to all vertices after
                                normalizing.
+        voxel_range (float): Range of voxelization.
 
     """
 
     def __init__(self, resolution: int,
                  normalize: bool = True,
-                 vertex_offset: float = 0.):
+                 vertex_offset: float = 0.,
+                 voxel_range: float = 1.):
         self.resolution = resolution
         self.normalize = normalize
         self.vertex_offset = vertex_offset
+        self.voxel_range = voxel_range
 
     def __call__(self, mesh: TriangleMesh):
         """
@@ -684,7 +687,8 @@ def __call__(self, mesh: TriangleMesh):
         """
         voxels = cvt.trianglemesh_to_voxelgrid(mesh, self.resolution,
                                                normalize=self.normalize,
-                                               vertex_offset=self.vertex_offset)
+                                               vertex_offset=self.vertex_offset,
+                                               voxel_range=self.voxel_range)
         return voxels