def __init__(self,
usd_filepath: str,
cache_dir: str = '../data/USDMeshes'):
usd_filepath = Path(usd_filepath)
assert usd_filepath.suffix in ['.usd', '.usda']
assert usd_filepath.exists(
), f'USD file at {usd_filepath} was not found.'
self.cache = helpers.Cache(get_mesh_attributes,
cache_dir,
cache_key=helpers._get_hash(usd_filepath))
self.names = self.cache.cached_ids
stage = Usd.Stage.Open(str(usd_filepath))
mesh_prims = [x for x in stage.Traverse() if UsdGeom.Mesh(x)]
uncached_mesh_prims = filter(lambda x: x.GetName() not in self.names,
mesh_prims)
for mesh_prim in uncached_mesh_prims:
name = mesh_prim.GetName()
mesh = UsdGeom.Mesh(mesh_prim)
face_counts = torch.tensor(mesh.GetFaceVertexCountsAttr().Get())
if not torch.allclose(face_counts, face_counts[0]):
log.warn(f'Skipping mesh {name}, not all faces have the same '
'number of vertices.')
else:
self.cache(name, usd_mesh=mesh)
def __init__(self, root: str, cache_dir: str, categories: list = ['chair'], train: bool = True,
split: float = .7, resolution: int = 100, smoothing_iterations: int = 3, mode='Tri',
no_progress: bool = False):
assert mode in ['Tri', 'Quad']
self.root = Path(root)
self.cache_dir = Path(cache_dir) / 'surface_meshes'
dataset_params = {
'root': root,
'categories': categories,
'train': train,
'split': split,
'no_progress': no_progress,
}
self.params = {
'resolution': resolution,
'smoothing_iterations': smoothing_iterations,
'mode': mode,
}
mesh_dataset = ShapeNet_Meshes(**dataset_params)
voxel_dataset = ShapeNet_Voxels(**dataset_params, cache_dir=cache_dir, resolutions=[resolution])
combined_dataset = ShapeNet_Combination([mesh_dataset, voxel_dataset])
self.names = combined_dataset.names
self.synset_idxs = combined_dataset.synset_idxs
self.synsets = combined_dataset.synsets
self.labels = combined_dataset.labels
if mode == 'Tri':
mesh_conversion = tfs.VoxelGridToTriangleMesh(threshold=0.5,
mode='marching_cubes',
normalize=False)
else:
mesh_conversion = tfs.VoxelGridToQuadMesh(threshold=0.5,
normalize=False)
def convert(og_mesh, voxel):
transforms = tfs.Compose([mesh_conversion,
tfs.MeshLaplacianSmoothing(smoothing_iterations)])
new_mesh = transforms(voxel)
new_mesh.vertices = pcfunc.realign(new_mesh.vertices, og_mesh.vertices)
return {'vertices': new_mesh.vertices, 'faces': new_mesh.faces}
self.cache_convert = helpers.Cache(convert, self.cache_dir,
cache_key=helpers._get_hash(self.params))
desc = 'converting to surface meshes'
for idx in tqdm(range(len(combined_dataset)), desc=desc, disable=no_progress):
name = combined_dataset.names[idx]
if name not in self.cache_convert.cached_ids:
sample = combined_dataset[idx]
voxel = sample['data'][str(resolution)]
og_mesh = TriangleMesh.from_tensors(sample['data']['vertices'],
sample['data']['faces'])
self.cache_convert(name, og_mesh=og_mesh, voxel=voxel)
def __init__(self, *args, preprocessing_transform=None, preprocessing_params: dict = None,
transform=None, no_progress: bool = False, **kwargs):
"""
Args:
positional and keyword arguments for initialize(*args, **kwargs) (see class and initialize documentation)
preprocessing_params (dict): parameters for the preprocessing:
- 'cache_dir': path to the cached preprocessed data.
- 'num_workers': number of process used in parallel for preprocessing (default: number of cores)
preprocessing_transform (Callable): Called on the outputs of _get_data over the indices
from 0 to len(self) during the construction of the dataset,
the preprocessed outputs are then cached to 'cache_dir'.
transform (Callable): Called on the preprocessed data at __getitem__.
no_progress (bool): disable tqdm progress bar for preprocessing.
"""
self.initialize(*args, **kwargs)
if preprocessing_transform is not None:
desc = 'Applying preprocessing'
if preprocessing_params is None:
preprocessing_params = {}
cache_dir = preprocessing_params.get('cache_dir')
assert cache_dir is not None, 'Cache directory is not given'
self.cache_convert = helpers.Cache(
preprocessing_transform,
cache_dir=cache_dir,
cache_key=helpers._get_hash(repr(preprocessing_transform))
)
use_cuda = preprocessing_params.get('use_cuda', False)
num_workers = preprocessing_params.get('num_workers')
uncached = [idx for idx in range(len(self)) if self._get_attributes(idx)[
'name'] not in self.cache_convert.cached_ids]
if len(uncached) > 0:
if num_workers == 0:
with torch.no_grad():
for idx in tqdm(range(len(self)), desc=desc, disable=no_progress):
name = self._get_attributes(idx)['name']
data = self._get_data(idx)
self.cache_convert(name, *data)
else:
p = Pool(num_workers)
iterator = p.imap_unordered(
_preprocess_task,
[(idx, self._get_data, self._get_attributes, self.cache_convert)
for idx in uncached])
for i in tqdm(range(len(uncached)), desc=desc, disable=no_progress):
next(iterator)
else:
self.cache_convert = None
self.transform = transform
def __init__(self, root: str, cache_dir: str, categories: list = ['chair'], train: bool = True,
split: float = .7, resolution: int = 100, num_points: int = 5000, occ: bool = False,
smoothing_iterations: int = 3, sample_box=True, no_progress: bool = False):
self.root = Path(root)
self.cache_dir = Path(cache_dir) / 'sdf_points'
self.params = {
'resolution': resolution,
'num_points': num_points,
'occ': occ,
'smoothing_iterations': smoothing_iterations,
'sample_box': sample_box,
}
surface_mesh_dataset = ShapeNet_Surface_Meshes(root=root,
cache_dir=cache_dir,
categories=categories,
train=train,
split=split,
resolution=resolution,
smoothing_iterations=smoothing_iterations,
no_progress=no_progress)
self.names = surface_mesh_dataset.names
self.synset_idxs = surface_mesh_dataset.synset_idxs
self.synsets = surface_mesh_dataset.synsets
self.labels = surface_mesh_dataset.labels
def convert(mesh):
sdf = mesh_cvt.trianglemesh_to_sdf(mesh, num_points)
bbox_true = torch.stack((mesh.vertices.min(dim=0)[0],
mesh.vertices.max(dim=0)[0]), dim=1).view(-1)
points = 1.05 * (torch.rand(self.params['num_points'], 3).to(mesh.vertices.device) - .5)
distances = sdf(points)
return {'points': points, 'distances': distances, 'bbox': bbox_true}
self.cache_convert = helpers.Cache(convert, self.cache_dir,
cache_key=helpers._get_hash(self.params))
desc = 'converting to sdf points'
for idx in tqdm(range(len(surface_mesh_dataset)), desc=desc, disable=no_progress):
name = surface_mesh_dataset.names[idx]
if name not in self.cache_convert.cached_ids:
idx = surface_mesh_dataset.names.index(name)
sample = surface_mesh_dataset[idx]
mesh = TriangleMesh.from_tensors(sample['data']['vertices'],
sample['data']['faces'])
# Use cuda if available to speed up conversion
if torch.cuda.is_available():
mesh.cuda()
self.cache_convert(name, mesh=mesh)
def __init__(self, root: str, cache_dir: str, categories: list = ['chair'], train: bool = True,
split: float = .7, num_points: int = 5000, smoothing_iterations=3,
surface=True, resolution=100, normals=True, no_progress: bool = False):
self.root = Path(root)
self.cache_dir = Path(cache_dir) / 'points'
dataset_params = {
'root': root,
'categories': categories,
'train': train,
'split': split,
'no_progress': no_progress,
}
self.params = {
'num_points': num_points,
'smoothing_iterations': smoothing_iterations,
'surface': surface,
'resolution': resolution,
'normals': normals,
}
if surface:
dataset = ShapeNet_Surface_Meshes(**dataset_params,
cache_dir=cache_dir,
resolution=resolution,
smoothing_iterations=smoothing_iterations)
else:
dataset = ShapeNet_Meshes(**dataset_params)
self.names = dataset.names
self.synset_idxs = dataset.synset_idxs
self.synsets = dataset.synsets
self.labels = dataset.labels
def convert(mesh):
points, face_choices = mesh_cvt.trianglemesh_to_pointcloud(mesh, num_points)
face_normals = mesh.compute_face_normals()
point_normals = face_normals[face_choices]
return {'points': points, 'normals': point_normals}
self.cache_convert = helpers.Cache(convert, self.cache_dir,
cache_key=helpers._get_hash(self.params))
desc = 'converting to points'
for idx in tqdm(range(len(dataset)), desc=desc, disable=no_progress):
name = dataset.names[idx]
if name not in self.cache_convert.cached_ids:
idx = dataset.names.index(name)
sample = dataset[idx]
mesh = TriangleMesh.from_tensors(sample['data']['vertices'],
sample['data']['faces'])
self.cache_convert(name, mesh=mesh)
def test_cache_tensor(device):
tensor = torch.ones(5, device=device)
cache = helpers.Cache(func=lambda x: x,
cache_dir=CACHE_DIR,
cache_key='test')
cache('tensor', x=tensor)
# Make sure cache is created
assert os.path.exists(os.path.join(CACHE_DIR, 'test', 'tensor.p'))
# Confirm loaded tensor is correct and on CPU device
loaded = cache('tensor')
assert torch.all(loaded.eq(tensor.cpu()))
def test_cache_pointcloud(device):
points = torch.ones(10, 3, device=device)
pointcloud = kal.rep.PointCloud(points)
cache = helpers.Cache(func=lambda x: x,
cache_dir=CACHE_DIR,
cache_key='test')
cache('pointcloud', x=pointcloud)
# Make sure cache is created
assert os.path.exists(os.path.join(CACHE_DIR, 'test', 'pointcloud.p'))
# Confirm loaded pointcloud is correct and on CPU device
loaded = cache('pointcloud')
assert torch.all(loaded.points.eq(points.cpu()))
def test_cache_voxelgrid(device):
voxels = torch.ones(3, 3, 3, device=device)
voxelgrid = kal.rep.VoxelGrid(voxels)
cache = helpers.Cache(func=lambda x: x,
cache_dir=CACHE_DIR,
cache_key='test')
cache('voxelgrid', x=voxelgrid)
# Make sure cache is created
assert os.path.exists(os.path.join(CACHE_DIR, 'test', 'voxelgrid.p'))
# Confirm loaded voxelgrid is correct and on CPU device
loaded = cache('voxelgrid')
assert torch.all(loaded.voxels.eq(voxels.cpu()))
def test_cache_mesh(device):
vertices = torch.ones(10, 3, device=device)
faces = torch.ones(20, 3, device=device, dtype=torch.long)
mesh = kal.rep.TriangleMesh.from_tensors(vertices, faces)
cache = helpers.Cache(func=lambda x: x,
cache_dir=CACHE_DIR,
cache_key='test')
cache('mesh', x=mesh)
# Make sure cache is created
assert os.path.exists(os.path.join(CACHE_DIR, 'test', 'mesh.p'))
# Confirm loaded mesh is correct and on CPU device
loaded = cache('mesh')
assert torch.all(loaded.vertices.eq(vertices.cpu()))
assert torch.all(loaded.faces.eq(faces.cpu()))
def test_cache_dict(device):
dictionary = {
'a': torch.ones(5, device=device),
'b': np.zeros(5),
}
cache = helpers.Cache(func=lambda x: x,
cache_dir=CACHE_DIR,
cache_key='test')
cache('dictionary', x=dictionary)
# Make sure cache is created
assert os.path.exists(os.path.join(CACHE_DIR, 'test', 'dictionary.p'))
# Confirm loaded dict is correct and on CPU device
loaded = cache('dictionary')
assert torch.all(loaded['a'].eq(dictionary['a'].cpu()))
assert np.all(np.isclose(loaded['b'], dictionary['b']))
