Add voxel_range option for mesh to voxel conversion

NVIDIAGameWorks · Jul 10, 2020 · 413c655 · 413c655
1 parent ec193ff
commit 413c655
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 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
@@ -387,7 +387,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 +407,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)