diff --git a/kaolin/conversions/meshconversions.py b/kaolin/conversions/meshconversions.py index 9011db32e..f9ce99fa5 100644 --- 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 index 1646f9ce6..56de57af8 100644 --- 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 index a2162fe58..5a68b7aa3 100644 --- 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