def import_pointclouds(file_path, scene_paths=None, times=None):
r"""Import one or more pointclouds from a USD file.
Assumes that pointclouds are interpreted using point instancers. Converts the coordinates
of each point instance to a point within the output pointcloud.
Args:
file_path (str): Path to usd file (\*.usd, \*.usda).
scene_paths (list of str, optional): Scene path(s) within the USD file indicating which primitive(s)
to import. If None, will return all pointclouds found based on PointInstancer prims with `kaolin_type`
primvar set to `PointCloud`.
times (list of int): Positive integers indicating the time at which to retrieve parameters.
Returns:
(list of torch.FloatTensor): Point coordinates.
Example:
>>> points = torch.rand(100, 3)
>>> stage = export_pointclouds('./new_stage.usd', [points, points, points])
>>> pointclouds = import_pointclouds(file_path='./new_stage.usd')
>>> len(pointclouds)
3
>>> pointclouds[0].shape
torch.Size([100, 3])
"""
assert os.path.exists(file_path)
stage = Usd.Stage.Open(file_path)
# If scene path not specified, find all point clouds
if scene_paths is None:
scene_paths = []
for p in stage.Traverse():
is_point_instancer = UsdGeom.PointInstancer(p)
if UsdGeom.PointInstancer(p) and p.GetAttribute(
'primvars:kaolin_type').Get() == 'PointCloud':
scene_paths.append(p.GetPath())
if times is None:
times = [Usd.TimeCode.Default()] * len(scene_paths)
pointclouds = []
for scene_path, time in zip(scene_paths, times):
prim = stage.GetPrimAtPath(scene_path)
assert prim, f'The prim at {scene_path} does not exist.'
instancer = UsdGeom.PointInstancer(prim)
assert instancer # Currently only support pointclouds from point instancers
pointclouds.append(
torch.tensor(instancer.GetPositionsAttr().Get(time=time)))
return pointclouds
def testManipulatePointInstancePosition(self):
# Create a UFE path to a PointInstancer prim with an instanceIndex on
# the end. This path uniquely identifies a specific point instance.
# We also pick one with a non-zero initial position.
instanceIndex = 7
ufePath = ufe.Path([
mayaUtils.createUfePathSegment('|UsdProxy|UsdProxyShape'),
usdUtils.createUfePathSegment(
'/PointInstancerGrid/PointInstancer/%d' % instanceIndex)
])
ufeItem = ufe.Hierarchy.createItem(ufePath)
# Select the point instance scene item.
globalSelection = ufe.GlobalSelection.get()
globalSelection.append(ufeItem)
# Get the PointInstancer prim for validating the values in USD.
ufePathString = ufe.PathString.string(ufePath)
prim = mayaUsdUfe.ufePathToPrim(ufePathString)
pointInstancer = UsdGeom.PointInstancer(prim)
self.assertTrue(pointInstancer)
# The PointInstancer should have 14 authored positions initially.
positionsAttr = pointInstancer.GetPositionsAttr()
positions = positionsAttr.Get()
self.assertEqual(len(positions), 14)
# Validate the initial position before manipulating
position = positions[instanceIndex]
self.assertTrue(
Gf.IsClose(position, Gf.Vec3f(-4.5, 1.5, 0.0), self.EPSILON))
# Perfom a translate manipulation via the move command.
cmds.move(1.0, 2.0, 3.0, objectSpace=True, relative=True)
# Re-fetch the USD positions and check for the update.
position = positionsAttr.Get()[instanceIndex]
self.assertTrue(
Gf.IsClose(position, Gf.Vec3f(-3.5, 3.5, 3.0), self.EPSILON))
# Try another move.
cmds.move(6.0, 5.0, 4.0, objectSpace=True, relative=True)
# Re-fetch the USD positions and check for the update.
position = positionsAttr.Get()[instanceIndex]
self.assertTrue(
Gf.IsClose(position, Gf.Vec3f(2.5, 8.5, 7.0), self.EPSILON))
# Now undo, and re-check.
cmds.undo()
position = positionsAttr.Get()[instanceIndex]
self.assertTrue(
Gf.IsClose(position, Gf.Vec3f(-3.5, 3.5, 3.0), self.EPSILON))
# And once more.
cmds.undo()
position = positionsAttr.Get()[instanceIndex]
self.assertTrue(
Gf.IsClose(position, Gf.Vec3f(-4.5, 1.5, 0.0), self.EPSILON))
def test_OneInstanceProtoXform(self):
stage = Usd.Stage.Open("test.usda")
pi = UsdGeom.PointInstancer(
stage.GetPrimAtPath("/OneInstanceProtoXform"))
# Test with prototype xforms (default).
baseTime = 0
tr = timeRange(baseTime)
xformsArray = self.computeInstanceTransforms(pi, tr, baseTime)
compares = [[
Gf.Matrix4d(Gf.Rotation(Gf.Vec3d(0, 0, 1), 180 + time * 36),
Gf.Vec3d(time * 5, time * 10, time * 20))
] for time, delta in tr]
self.assertAllMatrixListsEqual(xformsArray, compares)
# Test without prototype xforms.
baseTime = 0
tr = timeRange(baseTime)
xformsArray = self.computeInstanceTransforms(
pi, tr, baseTime, UsdGeom.PointInstancer.ExcludeProtoXform)
compares = [[
Gf.Matrix4d(Gf.Rotation(Gf.Vec3d(0, 0, 1), time * 36),
Gf.Vec3d(time * 5, time * 10, time * 20))
] for time, delta in tr]
self.assertAllMatrixListsEqual(xformsArray, compares)
def test_ComputeInstancerCount(self):
stage = Usd.Stage.Open('instancer.usda')
unset = UsdGeom.PointInstancer.Get(stage, '/UnsetIndices')
blocked = UsdGeom.PointInstancer.Get(stage, '/BlockedIndices')
empty = UsdGeom.PointInstancer.Get(stage, '/EmptyIndices')
timeSampled = UsdGeom.PointInstancer.Get(stage, '/TimeSampledIndices')
timeSampledAndDefault = UsdGeom.PointInstancer.Get(stage, '/TimeSampledAndDefaultIndices')
testTimeSamples = [
(unset, Usd.TimeCode.EarliestTime(), 0),
(blocked, Usd.TimeCode.EarliestTime(), 0),
(empty, Usd.TimeCode.EarliestTime(), 0),
(timeSampled, Usd.TimeCode.EarliestTime(), 3),
(timeSampledAndDefault, Usd.TimeCode.EarliestTime(), 5)]
testDefaults = [
(unset, 0),
(blocked, 0),
(empty, 0),
(timeSampled, 0),
(timeSampledAndDefault, 4)]
for (schema, timeCode, expected) in testTimeSamples:
self.assertTrue(schema)
self.assertEqual(schema.GetInstanceCount(timeCode), expected)
for (schema, expected) in testDefaults:
self.assertTrue(schema)
self.assertEqual(schema.GetInstanceCount(), expected)
invalid = UsdGeom.PointInstancer(Usd.Prim())
self.assertFalse(invalid)
with self.assertRaises(RuntimeError):
self.assertEqual(invalid.GetInstanceCount(), 0)
with self.assertRaises(RuntimeError):
self.assertEqual(invalid.GetInstanceCount(
Usd.TimeCode.EarliestTime()), 0)
def test_NoInstancesDefault(self):
stage = Usd.Stage.Open("test.usda")
pi = UsdGeom.PointInstancer(stage.GetPrimAtPath("/NoInstances"))
xforms = pi.ComputeInstanceTransformsAtTime(Usd.TimeCode.Default(),
Usd.TimeCode.Default())
self.assertEqual(len(xforms), 0)
def test_OneInstanceNoSamplesDefault(self):
stage = Usd.Stage.Open("test.usda")
pi = UsdGeom.PointInstancer(
stage.GetPrimAtPath("/OneInstanceNoSamples"))
xformsArray = pi.ComputeInstanceTransformsAtTimes(
[Usd.TimeCode.Default()], Usd.TimeCode.Default())
compares = [[Gf.Matrix4d(1)]]
self.assertAllMatrixListsEqual(xformsArray, compares)
def test_NoInstances(self):
stage = Usd.Stage.Open("test.usda")
pi = UsdGeom.PointInstancer(stage.GetPrimAtPath("/NoInstances"))
for baseTime, _ in timeRange(0):
tr = timeRange(baseTime)
xformsArray = self.computeInstanceTransforms(pi, tr, baseTime)
for xforms in xformsArray:
self.assertEqual(len(xforms), 0)
def add_pointcloud(stage, points, scene_path, time=None):
r"""Add a pointcloud to an existing USD stage.
Create a pointcloud represented by point instances of a sphere centered at each point coordinate.
The stage is modified but not saved to disk.
Args:
stage (Usd.Stage): Stage onto which to add the pointcloud.
points (torch.FloatTensor): Pointcloud tensor containing ``N`` points of shape ``(N, 3)``.
scene_path (str): Absolute path of pointcloud within the USD file scene. Must be a valid Sdf.Path.
time (int, optional): Positive integer defining the time at which the supplied parameters correspond to.
Returns:
(Usd.Stage)
Example:
>>> stage = create_stage('./new_stage.usd')
>>> points = torch.rand(100, 3)
>>> stage = add_pointcloud(stage, points, '/World/PointClouds/pointcloud_0')
>>> stage.Save()
"""
scene_path = Sdf.Path(scene_path)
if time is None:
time = Usd.TimeCode.Default()
if stage.GetPrimAtPath(scene_path):
instancer_prim = stage.GetPrimAtPath(scene_path)
else:
instancer_prim = stage.DefinePrim(scene_path, 'PointInstancer')
instancer = UsdGeom.PointInstancer(instancer_prim)
assert instancer
sphere = UsdGeom.Sphere.Define(stage, f'{scene_path}/sphere')
sphere.GetRadiusAttr().Set(0.5)
instancer.CreatePrototypesRel().SetTargets([sphere.GetPath()])
# Calculate default point scale
bounds = points.max(dim=0)[0] - points.min(dim=0)[0]
min_bound = min(bounds)
scale = (min_bound / points.size(0)**(1 / 3)).item()
# Generate instancer parameters
indices = [0] * points.size(0)
positions = points.cpu().tolist()
scales = [(scale, ) * 3] * points.size(0)
# Populate point instancer
instancer.GetProtoIndicesAttr().Set(indices, time=time)
instancer.GetPositionsAttr().Set(positions, time=time)
instancer.GetScalesAttr().Set(scales, time=time)
# Create a primvar to identify the point instancer as a Kaolin PointCloud
prim = stage.GetPrimAtPath(instancer.GetPath())
pv = UsdGeom.PrimvarsAPI(prim).CreatePrimvar('kaolin_type',
Sdf.ValueTypeNames.String)
pv.Set('PointCloud')
return stage
def test_OneInstanceAcceleration(self):
stage = Usd.Stage.Open("test.usda")
pi = UsdGeom.PointInstancer(
stage.GetPrimAtPath("/OneInstanceAcceleration"))
# Test directly on sample.
baseTime = 0
tr = timeRange(baseTime)
xformsArray = self.computeInstanceTransforms(pi, tr, baseTime)
compares = [[
Gf.Matrix4d(
Gf.Rotation(Gf.Vec3d(0, 0, 1), time * 36),
Gf.Vec3d((time / 24.0) * (120 + (time * 1 * 0.5)),
(time / 24.0) * (120 + (time * 1 * 0.5)),
(time / 24.0) * (120 + (time * 1 * 0.5))))
] for time, delta in tr]
self.assertAllMatrixListsEqual(xformsArray, compares)
# Test in-between samples.
baseTime = 2
tr = timeRange(baseTime)
xformsArray = self.computeInstanceTransforms(pi, tr, baseTime)
compares = [[
Gf.Matrix4d(
Gf.Rotation(Gf.Vec3d(0, 0, 1), time * 36),
Gf.Vec3d((time / 24.0) * (120 + (time * 1 * 0.5)),
(time / 24.0) * (120 + (time * 1 * 0.5)),
(time / 24.0) * (120 + (time * 1 * 0.5))))
] for time, delta in tr]
self.assertAllMatrixListsEqual(xformsArray, compares)
# Test with basetime before natural sample.
baseTime = 5
tr = timeRange(baseTime)
xformsArray = self.computeInstanceTransforms(pi, tr, baseTime - 1)
compares = [[
Gf.Matrix4d(
Gf.Rotation(Gf.Vec3d(0, 0, 1), time * 36),
Gf.Vec3d((time / 24.0) * (120 + (time * 1 * 0.5)),
(time / 24.0) * (120 + (time * 1 * 0.5)),
(time / 24.0) * (120 + (time * 1 * 0.5))))
] for time, delta in tr]
self.assertAllMatrixListsEqual(xformsArray, compares)
# Test with basetime on natural sample.
baseTime = 5
tr = timeRange(baseTime)
xformsArray = self.computeInstanceTransforms(pi, tr, baseTime)
compares = [[
Gf.Matrix4d(
Gf.Rotation(Gf.Vec3d(0, 0, 1), 180 - delta * 36),
Gf.Vec3d(25 + (delta / 24.0) * (120 + (delta * 1 * 0.5)),
50 + (delta / 24.0) * (240 + (delta * 2 * 0.5)),
100 + (delta / 24.0) * (480 + (delta * 3 * 0.5))))
] for time, delta in tr]
self.assertAllMatrixListsEqual(xformsArray, compares)
def test_OneInstanceNoSamples(self):
stage = Usd.Stage.Open("test.usda")
pi = UsdGeom.PointInstancer(
stage.GetPrimAtPath("/OneInstanceNoSamples"))
baseTime = 1
tr = timeRange(baseTime)
xformsArray = self.computeInstanceTransforms(pi, tr, baseTime)
compares = [[Gf.Matrix4d(1)] for time, delta in tr]
self.assertAllMatrixListsEqual(xformsArray, compares)
def doTest(self, namespacedInstancer, namespacedPrototype,
stripNamespaces):
cmds.file(self.mayaFile, open=True, force=True)
group = 'group1'
namespace = 'mayaNamespace'
instancer = baseInstancer = 'instancer1'
cubeTrans = baseCubeTrans = 'pCube1'
cubeShape = baseCubeShape = 'pCubeShape1'
# Note: the scene already has the namespace defined, so we don't
# need to do that
if namespacedInstancer:
instancer = '{}:{}'.format(namespace, baseInstancer)
cmds.rename(baseInstancer, instancer)
self.assertFalse(cmds.objExists(baseInstancer))
self.assertTrue(cmds.objExists(instancer))
if namespacedPrototype:
cubeTrans = '{}:{}'.format(namespace, baseCubeTrans)
cubeShape = '{}:{}'.format(namespace, baseCubeShape)
# renaming trans should also rename shape
cmds.rename(baseCubeTrans, cubeTrans)
self.assertFalse(cmds.objExists(baseCubeTrans))
self.assertFalse(cmds.objExists(baseCubeShape))
self.assertTrue(cmds.objExists(cubeTrans))
self.assertTrue(cmds.objExists(cubeShape))
stage = self.doExport(stripNamespaces=stripNamespaces)
self.assertTrue(stage)
if namespacedInstancer and not stripNamespaces:
instancerPrimPath = '/{}/{}_{}'.format(group, namespace,
baseInstancer)
else:
instancerPrimPath = '/{}/{}'.format(group, baseInstancer)
instancerPrim = stage.GetPrimAtPath(instancerPrimPath)
self.assertTrue(instancerPrim)
instancer = UsdGeom.PointInstancer(instancerPrim)
self.assertTrue(instancer)
if namespacedPrototype and not stripNamespaces:
protoPrimPath = '{}/Prototypes/{}_{}_0'.format(
instancerPrimPath, namespace, baseCubeTrans)
else:
protoPrimPath = '{}/Prototypes/{}_0'.format(
instancerPrimPath, baseCubeTrans)
protoPrim = stage.GetPrimAtPath(protoPrimPath)
self.assertTrue(protoPrim)
protoMesh = UsdGeom.Mesh(protoPrim)
self.assertTrue(protoMesh)
prototypes = instancer.GetPrototypesRel().GetTargets()
self.assertEqual(prototypes, [protoPrim.GetPath()])
def testMashPrototypes_NoIdsArray(self):
"""
MASH instancers might not have an ids array if using dynamics.
Make sure they still export OK.
"""
instancerPrim = self.stage.GetPrimAtPath(
"/InstancerTest/MASH2_Instancer")
self.assertTrue(instancerPrim)
instancer = UsdGeom.PointInstancer(instancerPrim)
protoIndices = instancer.GetProtoIndicesAttr().Get()
self.assertEqual(len(protoIndices), 10)
def testTransforms(self):
"""
Check that the point transforms are correct.
"""
mayaInstancer = OMFX.MFnInstancer(self._GetDagPath("instancer1"))
usdInstancer = UsdGeom.PointInstancer(
self.stage.GetPrimAtPath("/InstancerTest/instancer1"))
time = self.START_TIMECODE
while time <= self.END_TIMECODE:
cmds.currentTime(time, edit=True)
# Need to do this because MFnInstancer will give instance matrices
# as offsets from prototypes' original world space positions.
worldPositions = [
self._GetWorldSpacePosition("|dummyGroup|pCube1"),
self._GetWorldSpacePosition(
"|InstancerTest|instancer1|prototypeUnderInstancer"),
self._GetWorldSpacePosition("|referencePrototype")
]
paths = OM.MDagPathArray()
matrices = OM.MMatrixArray()
particlePathStartIndices = OM.MIntArray()
pathIndices = OM.MIntArray()
mayaInstancer.allInstances(paths, matrices,
particlePathStartIndices, pathIndices)
usdInstanceTransforms = \
usdInstancer.ComputeInstanceTransformsAtTime(time, time)
usdProtoIndices = usdInstancer.GetProtoIndicesAttr().Get(time)
self.assertEqual(matrices.length(), len(usdInstanceTransforms))
# Compute the instancer-space position of instances in Maya
# (including the protos' transforms). By default, this is what
# UsdGeomPointInstancer::ComputeInstanceTransformsAtTime already
# gives us.
mayaWorldPositions = [
worldPositions[protoIndex] for protoIndex in usdProtoIndices
]
mayaGfMatrices = [
mayaWorldPositions[i] * self._MayaToGfMatrix(matrices[i])
for i in xrange(matrices.length())
]
usdGfMatrices = [
usdInstanceTransforms[i]
for i in xrange(len(usdInstanceTransforms))
]
for i in xrange(len(usdGfMatrices)):
self._AssertXformMatrices(mayaGfMatrices[i], usdGfMatrices[i])
time += 1.0
def GetInstanceIndicesForIds(prim, instanceIds, time):
'''Attempt to find the instance indices of a list of authored instance IDs
for prim 'prim' at time 'time'. If the prim is not a PointInstancer or does
not have authored IDs, returns None. If any ID from 'instanceIds' does not
exist at the given time, its index is not added to the list (because it does
not have an index).'''
ids = UsdGeom.PointInstancer(prim).GetIdsAttr().Get(time)
if ids:
return [instanceIndex for instanceIndex, instanceId in enumerate(ids)
if instanceId in instanceIds]
else:
return None
def test_Extent(self):
stage = Usd.Stage.Open("test.usda")
pi = UsdGeom.PointInstancer(
stage.GetPrimAtPath("/MultiInstanceForExtents"))
times = [0, 1, 2]
expectedExtents = [[(-1, -1, -1), (1, 1, 1)],
[(-3.7600734, 1.2399265, -1),
(3.7600734, 6.2600737, 3.5)],
[(-6.3968024, 3.6031978, -1),
(6.3968024, 11.396802, 6)]]
self.compareExtents(pi, times, 0, expectedExtents)
def GetInstanceIdForIndex(prim, instanceIndex, time):
'''Attempt to find an authored Id value for the instance at index
'instanceIndex' at time 'time', on the given prim 'prim', which we access
as a UsdGeom.PointInstancer (whether it actually is or not, to provide
some dynamic duck-typing for custom instancer types that support Ids.
Returns 'None' if no ids attribute was found, or if instanceIndex is
outside the bounds of the ids array.'''
if not prim or instanceIndex < 0:
return None
ids = UsdGeom.PointInstancer(prim).GetIdsAttr().Get(time)
if not ids or instanceIndex >= len(ids):
return None
return ids[instanceIndex]
def get_pointcloud_bracketing_time_samples(stage, scene_path, target_time):
"""Returns two time samples that bracket target_time for point cloud attributes at a specified
scene_path.
Args:
stage (Usd.Stage)
scene_path (str)
target_time (Number)
Returns:
(iterable of 2 numbers)
"""
# Note: can also get usd_attr.GetTimeSamples()
prim = stage.GetPrimAtPath(scene_path)
if UsdGeom.Points(prim):
geom_points = UsdGeom.Points(prim)
result = geom_points.GetPointsAttr().GetBracketingTimeSamples(target_time)
elif UsdGeom.PointInstancer(prim):
instancer = UsdGeom.PointInstancer(prim)
result = instancer.GetPositionsAttr().GetBracketingTimeSamples(target_time)
else:
raise TypeError("The prim is neither UsdGeomPoints nor UsdGeomPointInstancer.")
return result
def test_OneInstanceVelocityScale(self):
stage = Usd.Stage.Open("test.usda")
pi = UsdGeom.PointInstancer(
stage.GetPrimAtPath("/OneInstanceVelocityScale"))
# Test when the velocityScale is set to 2.
baseTime = 0
tr = timeRange(baseTime)
xformsArray = self.computeInstanceTransforms(pi, tr, baseTime)
compares = [[
Gf.Matrix4d(Gf.Rotation(Gf.Vec3d(0, 0, 1), time * 72),
Gf.Vec3d(time * 10, time * 20, time * 40))
] for time, delta in tr]
self.assertAllMatrixListsEqual(xformsArray, compares)
def test_Mask(self):
stage = Usd.Stage.Open("test.usda")
pi = UsdGeom.PointInstancer(stage.GetPrimAtPath("/MultiInstanceMask"))
# Test with 3 instances with the second masked out.
baseTime = 0
tr = timeRange(baseTime)
xformsArray = self.computeInstanceTransforms(pi, tr, baseTime)
compares = [[
Gf.Matrix4d(Gf.Rotation(Gf.Vec3d(0, 0, 1), time * 36),
Gf.Vec3d(time * 5, time * 10, time * 20)),
Gf.Matrix4d(Gf.Rotation(Gf.Vec3d(0, 0, 1), 180 - time * 36),
Gf.Vec3d(time * 5, time * 10, 2 + time * 20))
] for time, delta in tr]
self.assertAllMatrixListsEqual(xformsArray, compares)
def get_pointcloud_bracketing_time_samples(stage, scene_path, target_time):
"""Returns two time samples that bracket target_time for point cloud attributes at a specified
scene_path.
Args:
stage (Usd.Stage)
scene_path (str)
target_time (Number)
Returns:
(iterable of 2 numbers)
"""
# Note: can also get usd_attr.GetTimeSamples()
prim = stage.GetPrimAtPath(scene_path)
instancer = UsdGeom.PointInstancer(prim)
assert instancer, f'Only point clouds from point instancers are supported. No PointInstancer at {scene_path}.'
return instancer.GetPositionsAttr().GetBracketingTimeSamples(target_time)
def testInstancePaths(self):
"""
Checks that the proto index assigned for each point is correct.
"""
mayaInstancer = OMFX.MFnInstancer(self._GetDagPath("instancer1"))
usdInstancer = UsdGeom.PointInstancer(
self.stage.GetPrimAtPath("/InstancerTest/instancer1"))
time = self.START_TIMECODE
while time <= self.END_TIMECODE:
cmds.currentTime(time, edit=True)
paths = OM.MDagPathArray()
matrices = OM.MMatrixArray()
particlePathStartIndices = OM.MIntArray()
pathIndices = OM.MIntArray()
mayaInstancer.allInstances(paths, matrices,
particlePathStartIndices, pathIndices)
usdProtoIndices = usdInstancer.GetProtoIndicesAttr().Get(time)
# Mapping of proto index to index(es) in the paths array.
# Note that in the Maya instancer a single point may map to multiple
# DAG paths, which correspond to all the shapes in the instanced
# hierarchy.
usdIndicesToMayaIndices = {
0: [0, 1], # the first prototype has two shapes in hierarchy
1: [2], # this prototype only has one shape
2: [3], # the reference prototype only has one shape
}
for i in xrange(len(usdProtoIndices)):
usdProtoIndex = usdProtoIndices[i]
expectedMayaIndices = usdIndicesToMayaIndices[usdProtoIndex]
mayaIndicesStart = particlePathStartIndices[i]
mayaIndicesEnd = particlePathStartIndices[i + 1]
self.assertEqual(mayaIndicesEnd - mayaIndicesStart,
len(expectedMayaIndices))
actualPathIndices = [
pathIndices[i]
for i in xrange(mayaIndicesStart, mayaIndicesEnd)
]
self.assertEqual(actualPathIndices, expectedMayaIndices)
time += 1.0
def test_OneInstanceUnalignedData(self):
stage = Usd.Stage.Open("test.usda")
pi = UsdGeom.PointInstancer(
stage.GetPrimAtPath("/OneInstanceUnalignedData"))
# Test that unaligned positions/orientations are handled properly.
baseTime = 3
tr = timeRange(baseTime)
xformsArray = self.computeInstanceTransforms(pi, tr, baseTime)
compares = []
for time, delta in tr:
rotationTime = time - 2
velocityTime = time - 1
compares.append([
Gf.Matrix4d(
Gf.Rotation(Gf.Vec3d(0, 0, 1), rotationTime * 36),
Gf.Vec3d(velocityTime * 5, velocityTime * 10,
velocityTime * 20))
])
self.assertAllMatrixListsEqual(xformsArray, compares)
def test_OneInstanceNoVelocities(self):
stage = Usd.Stage.Open("test.usda")
pi = UsdGeom.PointInstancer(
stage.GetPrimAtPath("/OneInstanceNoVelocities"))
# Test directly on sample.
baseTime = 0
tr = timeRange(baseTime)
xformsArray = self.computeInstanceTransforms(pi, tr, baseTime)
compares = []
for time, delta in tr:
if time < 0:
# Samples at times less than 0 should clamp to first sample.
compare = [Gf.Matrix4d(1)]
else:
compare = [
Gf.Matrix4d(Gf.Rotation(Gf.Vec3d(0, 0, 1), time * 36),
Gf.Vec3d(time * 5, time * 10, time * 20))
]
compares.append(compare)
self.assertAllMatrixListsEqual(xformsArray, compares)
# Test in-between samples.
baseTime = 2
tr = timeRange(baseTime)
xformsArray = self.computeInstanceTransforms(pi, tr, baseTime)
compares = [[
Gf.Matrix4d(Gf.Rotation(Gf.Vec3d(0, 0, 1), time * 36),
Gf.Vec3d(time * 5, time * 10, time * 20))
] for time, delta in tr]
self.assertAllMatrixListsEqual(xformsArray, compares)
# Test with basetime before and after natural sample. Since we are
# interpolating, these should always be the same.
baseTime = 5
tr = timeRange(baseTime)
xformsArrayBefore = self.computeInstanceTransforms(
pi, tr, baseTime - 1)
xformsArrayAfter = self.computeInstanceTransforms(pi, tr, baseTime)
self.assertAllMatrixListsEqual(xformsArrayBefore, xformsArrayAfter)
def test_PointInstancer(self):
stage = Usd.Stage.Open("testPointInstancer.usda")
pi = UsdGeom.PointInstancer(stage.GetPrimAtPath("/Instancer"))
# Verify the extent computation when no transform matrix is given.
self.verifyExtent(pi, [(-1, -1, -1), (3.5, 3.5, 3.5)])
# Verify the extent computation when no transform matrix is given.
self.verifyExtent(pi, [(-1, -1, -1), (3.5, 3.5, 3.5)],
Gf.Matrix4d(1.0))
# Apply an arbitrary transform matrix. Note that this is different from
# the extent displayed in Usdview. If you open testPointInstancer.usda
# and select /StrangeInstancer, you will see the box is larger than it
# needs to be (it transforms the bounding box around the entire
# PointInstancer, not the boxes for each instance like
# ComputeExtentFromPlugins). If you select
# /StrangeInstancerComputedExtent, it will display the (tighter)
# computed extent used below.
self.verifyExtent(
pi, [(-1.3977774381637573, 0.3787655532360077, 0.689382791519165),
(5.12123441696167, 6.12123441696167, 3.560617208480835)],
STRANGE_TRANSFORM)
def testManipulatePointInstanceScale(self):
# Create a UFE path to a PointInstancer prim with an instanceIndex on
# the end. This path uniquely identifies a specific point instance.
instanceIndex = 7
ufePath = ufe.Path([
mayaUtils.createUfePathSegment('|UsdProxy|UsdProxyShape'),
usdUtils.createUfePathSegment(
'/PointInstancerGrid/PointInstancer/%d' % instanceIndex)
])
ufeItem = ufe.Hierarchy.createItem(ufePath)
# Select the point instance scene item.
globalSelection = ufe.GlobalSelection.get()
globalSelection.append(ufeItem)
# Get the PointInstancer prim for validating the values in USD.
ufePathString = ufe.PathString.string(ufePath)
prim = mayaUsdUfe.ufePathToPrim(ufePathString)
pointInstancer = UsdGeom.PointInstancer(prim)
self.assertTrue(pointInstancer)
# The PointInstancer should not have any authored scales initially.
scalesAttr = pointInstancer.GetScalesAttr()
self.assertFalse(scalesAttr.HasAuthoredValue())
# Perfom a scale manipulation via the scale command.
cmds.scale(1.0, 2.0, 3.0, objectSpace=True, relative=True)
# The initial scale should have filled out the scales attribute.
scales = scalesAttr.Get()
self.assertEqual(len(scales), 14)
# Validate the scaled item.
scale = scales[instanceIndex]
self.assertTrue(
Gf.IsClose(scale, Gf.Vec3f(1.0, 2.0, 3.0), self.EPSILON))
# The non-scaled items should all have identity scales.
for i in [idx for idx in range(14) if idx != instanceIndex]:
scale = scales[i]
self.assertTrue(
Gf.IsClose(scale, Gf.Vec3f(1.0, 1.0, 1.0), self.EPSILON))
# Perfom another scale.
cmds.scale(4.0, 5.0, 6.0, objectSpace=True, relative=True)
# Re-fetch the USD scale and check for the update.
scale = scalesAttr.Get()[instanceIndex]
self.assertTrue(
Gf.IsClose(scale, Gf.Vec3f(4.0, 10.0, 18.0), self.EPSILON))
# Now undo, and re-check.
cmds.undo()
scale = scalesAttr.Get()[instanceIndex]
self.assertTrue(
Gf.IsClose(scale, Gf.Vec3f(1.0, 2.0, 3.0), self.EPSILON))
# And once more.
# Note that with more complete undo support, this would ideally clear
# the authored scales attribute, returning it to its true original
# state. For now, we just validate that the scale value is back to its
# default of identity.
cmds.undo()
scale = scalesAttr.Get()[instanceIndex]
self.assertTrue(
Gf.IsClose(scale, Gf.Vec3f(1.0, 1.0, 1.0), self.EPSILON))
def import_pointclouds(file_path, scene_paths=None, times=None):
r"""Import one or more pointclouds from a USD file.
Assumes that pointclouds are interpreted using point instancers or UsdGeomPoints. Converts the coordinates
of each point instance to a point within the output pointcloud.
Args:
file_path (str): Path to usd file (\*.usd, \*.usda).
scene_paths (list of str, optional): Scene path(s) within the USD file indicating which primitive(s)
to import. If None, will return all pointclouds found based on PointInstancer or UsdGeomPoints prims with `kaolin_type`
primvar set to `PointCloud`.
times (list of int): Positive integers indicating the time at which to retrieve parameters.
Returns:
list of namedtuple of:
- **points** (list of torch.FloatTensor): of shape (num_points, 3)
- **colors** (list of torch.FloatTensor): of shape (num_points, 3)
- **normals** (list of torch.FloatTensor): of shape (num_points, 2)
Example:
>>> points = torch.rand(100, 3)
>>> stage = export_pointclouds('./new_stage.usd', [points, points, points])
>>> pointclouds = import_pointclouds(file_path='./new_stage.usd')[0]
>>> len(pointclouds)
3
>>> pointclouds[0].shape
torch.Size([100, 3])
"""
assert os.path.exists(file_path)
if scene_paths is None:
scene_paths = get_pointcloud_scene_paths(file_path)
if times is None:
times = [Usd.TimeCode.Default()] * len(scene_paths)
pointclouds = []
colors = []
normals = []
stage = Usd.Stage.Open(file_path)
for scene_path, time in zip(scene_paths, times):
prim = stage.GetPrimAtPath(scene_path)
assert prim, f'The prim at {scene_path} does not exist.'
if UsdGeom.Points(prim):
geom_points = UsdGeom.Points(prim)
pointclouds.append(torch.tensor(geom_points.GetPointsAttr().Get(time=time)))
color = geom_points.GetDisplayColorAttr().Get(time=time)
if color is None:
colors.append(color)
else:
colors.append(torch.tensor(color))
elif UsdGeom.PointInstancer(prim):
instancer = UsdGeom.PointInstancer(prim)
pointclouds.append(torch.tensor(instancer.GetPositionsAttr().Get(time=time)))
colors.append(None)
else:
raise TypeError("The prim is neither UsdGeomPoints nor UsdGeomPointInstancer.")
# TODO: place holders for normals for now
normals = [None] * len(colors)
params = [pointclouds, colors, normals]
return [pointcloud_return_type(p, c, n) for p, c, n in zip(*params)]
def import_voxelgrids(file_path, scene_paths=None, times=None):
r"""Import one or more voxelgrids from a USD file.
Assumes that the USD voxelgrid is defined by a point instancer. Converts the coordinates
of each point instance to an occupied voxel. The output grid size is determined from the `grid_size`
primvar. If not specified, grid size will be determined by the axis with the largest number of occupied
voxels. The output voxelgrid will be of shape ``[grid_size, grid_size, grid_size]``.
Args:
file_path (str): Path to usd file (\*.usd, \*.usda).
scene_paths (list of str, optional): Scene path(s) within the USD file indicating which PointInstancer
primitive(s) to import. If None, will return all pointclouds found based on PointInstancer
prims with `kaolin_type` primvar set to `VoxelGrid`.
times (list of int): Positive integers indicating the time at which to retrieve parameters.
Returns:
(list of torch.BoolTensor)
Example:
>>> voxelgrid_1 = torch.rand(32, 32, 32) > 0.5
>>> voxelgrid_2 = torch.rand(32, 32, 32) > 0.5
>>> stage = export_voxelgrids('./new_stage.usd', [voxelgrid_1, voxelgrid_2])
>>> voxelgrid_imp = import_voxelgrids(file_path='./new_stage.usd')
>>> len(voxelgrid_imp)
2
>>> voxelgrid_imp[0].shape
torch.Size([32, 32, 32])
"""
assert os.path.exists(file_path)
stage = Usd.Stage.Open(file_path)
# If scene path not specified, find all point clouds
if scene_paths is None:
scene_paths = []
for p in stage.Traverse():
is_point_instancer = UsdGeom.PointInstancer(p)
if UsdGeom.PointInstancer(p) and p.GetAttribute('primvars:kaolin_type').Get() == 'VoxelGrid':
scene_paths.append(p.GetPath())
if times is None:
times = [Usd.TimeCode.Default()] * len(scene_paths)
voxelgrids = []
for scene_path, time in zip(scene_paths, times):
prim = stage.GetPrimAtPath(scene_path)
assert prim, f'The prim at {scene_path} does not exist.'
instancer = UsdGeom.PointInstancer(prim)
assert instancer # Currently only support pointclouds from point instancers
voxel_indices = torch.tensor(instancer.GetPositionsAttr().Get(time=time), dtype=torch.long)
bounds = voxel_indices.max(dim=0)[0]
max_bound = bounds.max()
grid_size = prim.GetAttribute('primvars:grid_size').Get(time=time)
if grid_size is not None:
assert max_bound < grid_size
else:
grid_size = max_bound
voxelgrid = torch.zeros([grid_size, grid_size, grid_size], dtype=torch.bool)
voxelgrid[voxel_indices[:, 0], voxel_indices[:, 1], voxel_indices[:, 2]] = 1.
voxelgrids.append(voxelgrid)
return voxelgrids
def add_voxelgrid(stage, voxelgrid, scene_path, time=None):
r"""Add a voxelgrid to an existing USD stage.
Add a voxelgrid where occupied voxels are defined by non-zero values. The voxelgrid is
represented by point instances of a cube centered at each occupied index coordinate and scaled.
The stage is modified but not saved to disk.
Args:
stage (Usd.Stage): Stage onto which to add the voxelgrid.
voxelgrid (torch.BoolTensor): Binary voxelgrid of shape ``(N, N, N)``.
scene_path (str): Absolute path of voxelgrid within the USD file scene. Must be a valid Sdf.Path.
time (int, optional): Positive integer defining the time at which the supplied parameters correspond to.
Returns:
(Usd.Stage)
Example:
>>> stage = create_stage('./new_stage.usd')
>>> voxelgrid = torch.rand(32, 32, 32) > 0.5
>>> stage = add_voxelgrid(stage, voxelgrid, '/World/VoxelGrids/voxelgrid_0')
>>> stage.Save()
"""
scene_path = Sdf.Path(scene_path)
if time is None:
time = Usd.TimeCode.Default()
if stage.GetPrimAtPath(scene_path):
instancer_prim = stage.GetPrimAtPath(scene_path)
else:
instancer_prim = stage.DefinePrim(scene_path, 'PointInstancer')
instancer = UsdGeom.PointInstancer(instancer_prim)
assert instancer
cube = UsdGeom.Cube.Define(stage, f'{scene_path}/cube')
cube.GetSizeAttr().Set(1.0)
instancer.CreatePrototypesRel().SetTargets([cube.GetPath()])
# Get each occupied voxel's centre coordinate
points = torch.nonzero(voxelgrid, as_tuple=False).float()
# Generate instancer parameters
indices = [0] * points.shape[0]
positions = points.cpu().tolist()
scales = [(1.,) * 3] * points.size(0)
# Populate point instancer
instancer.GetProtoIndicesAttr().Set(indices, time=time)
instancer.GetPositionsAttr().Set(positions, time=time)
instancer.GetScalesAttr().Set(scales, time=time)
# Centre at (0, 0, 0) and scale points
adjusted_scale = 1. / (voxelgrid.size(0) - 1)
scale_op = instancer.GetPrim().GetAttribute('xformOp:scale')
if not scale_op:
scale_op = UsdGeom.Xformable(instancer).AddScaleOp()
scale_op.Set((adjusted_scale,) * 3, time=time)
translate_op = instancer.GetPrim().GetAttribute('xformOp:translate')
if not translate_op:
UsdGeom.Xformable(instancer).AddTranslateOp()
translate_op.Set((-0.5,) * 3, time=time)
# Create a primvar to record the voxelgrid size
prim = stage.GetPrimAtPath(instancer.GetPath())
pv = UsdGeom.PrimvarsAPI(prim).CreatePrimvar('grid_size', Sdf.ValueTypeNames.Int)
pv.Set(voxelgrid.size(0))
# Create a primvar to identify the poitn instancer as a Kaolin VoxelGrid
pv = UsdGeom.PrimvarsAPI(prim).CreatePrimvar('kaolin_type', Sdf.ValueTypeNames.String)
pv.Set('VoxelGrid')
return stage
def test_OneInstanceTimeSampleCorrespondenceValidation(self):
stage = Usd.Stage.Open("test.usda")
piDiffNumberPositionsAndVelocities = UsdGeom.PointInstancer(
stage.GetPrimAtPath(
"/OneInstanceDifferingNumberPositionsAndVelocities"))
piUnalignedPositionsAndVelocities = UsdGeom.PointInstancer(
stage.GetPrimAtPath("/OneInstanceUnalignedPositionsAndVelocities"))
piPositionsOnly = UsdGeom.PointInstancer(
stage.GetPrimAtPath("/OneInstanceUnalignedPositionsOnly"))
baseTime = 2
tr = timeRange(baseTime)
xformsArrayDiffNumberPositionsAndVelocities = self.computeInstanceTransforms(
piDiffNumberPositionsAndVelocities, tr, baseTime)
xformsArrayUnalignedPositionsAndVelocities = self.computeInstanceTransforms(
piUnalignedPositionsAndVelocities, tr, baseTime)
xformsArrayPositionsOnly = self.computeInstanceTransforms(
piPositionsOnly, tr, baseTime)
# Test that time sample correspondence validation works for positions
# and velocities by comparing the transformations of the point instancers
# with invalid time sample correspondence with the transformations of the
# point instancer that only has time samples for positions
self.assertAllMatrixListsEqual(
xformsArrayDiffNumberPositionsAndVelocities,
xformsArrayPositionsOnly)
self.assertAllMatrixListsEqual(
xformsArrayUnalignedPositionsAndVelocities,
xformsArrayPositionsOnly)
piDiffNumberOrientationsAndAngularVelocities = UsdGeom.PointInstancer(
stage.GetPrimAtPath(
"/OneInstanceDifferingNumberOrientationsAndAngularVelocities"))
piUnalignedOrientationsAndAngularVelocities = UsdGeom.PointInstancer(
stage.GetPrimAtPath(
"/OneInstanceUnalignedOrientationsAndAngularVelocities"))
piOrientationsOnly = UsdGeom.PointInstancer(
stage.GetPrimAtPath("/OneInstanceUnalignedOrientationsOnly"))
baseTime = 2
tr = timeRange(baseTime)
xformsArrayDiffNumberOrientationsAndAngularVelocities = self.computeInstanceTransforms(
piDiffNumberOrientationsAndAngularVelocities, tr, baseTime)
xformsArrayUnalignedOrientationsAndAngularVelocities = self.computeInstanceTransforms(
piUnalignedOrientationsAndAngularVelocities, tr, baseTime)
xformsArrayOrientationsOnly = self.computeInstanceTransforms(
piOrientationsOnly, tr, baseTime)
# Test that time sample correspondence validation works for orientations
# and angular velocities by comparing the transformations of the point
# instancers with invalid time sample correspondence with the transformations
# of the point instancer that only has time samples for orientations
self.assertAllMatrixListsEqual(
xformsArrayDiffNumberOrientationsAndAngularVelocities,
xformsArrayOrientationsOnly)
self.assertAllMatrixListsEqual(
xformsArrayUnalignedOrientationsAndAngularVelocities,
xformsArrayOrientationsOnly)
piDiffNumberVelocitiesAndAccelerations = UsdGeom.PointInstancer(
stage.GetPrimAtPath(
"/OneInstanceDiffNumberVelocitiesAndAccelerations"))
piUnalignedVelocitiesAndAccelerations = UsdGeom.PointInstancer(
stage.GetPrimAtPath(
"/OneInstanceUnalignedVelocitiesAndAccelerations"))
piDiffNumberPositionsAndVelocitiesAndAccelerations = UsdGeom.PointInstancer(
stage.GetPrimAtPath(
"/OneInstanceDiffNumberPositionsAndVelocitiesAndAccelerations")
)
piPositionsAndVelocitiesOnly = UsdGeom.PointInstancer(
stage.GetPrimAtPath("/OneInstancePositionsAndVelocitiesOnly"))
baseTime = 2
tr = timeRange(baseTime)
xformsArrayDiffNumberVelocitiesAndAccelerations = self.computeInstanceTransforms(
piDiffNumberVelocitiesAndAccelerations, tr, baseTime)
xformsArrayUnalignedVelocitiesAndAccelerations = self.computeInstanceTransforms(
piUnalignedVelocitiesAndAccelerations, tr, baseTime)
xformsArrayDiffNumberPositionsAndVelocitiesAndAccelerations = self.computeInstanceTransforms(
piDiffNumberPositionsAndVelocitiesAndAccelerations, tr, baseTime)
xformsArrayPositionsAndVelocitiesOnly = self.computeInstanceTransforms(
piPositionsAndVelocitiesOnly, tr, baseTime)
# Test that time sample correspondence validation works for velocities
# and accelerations by comparing the transformations of the point instancers
# with invalid time sample correspondence with the transformations of the
# point instancer that only has time samples for positions and velocities
self.assertAllMatrixListsEqual(
xformsArrayDiffNumberVelocitiesAndAccelerations,
xformsArrayPositionsAndVelocitiesOnly)
self.assertAllMatrixListsEqual(
xformsArrayUnalignedVelocitiesAndAccelerations,
xformsArrayPositionsAndVelocitiesOnly)
self.assertAllMatrixListsEqual(
xformsArrayDiffNumberPositionsAndVelocitiesAndAccelerations,
xformsArrayPositionsOnly)
def testManipulatePointInstanceOrientation(self):
# Create a UFE path to a PointInstancer prim with an instanceIndex on
# the end. This path uniquely identifies a specific point instance.
instanceIndex = 7
ufePath = ufe.Path([
mayaUtils.createUfePathSegment('|UsdProxy|UsdProxyShape'),
usdUtils.createUfePathSegment(
'/PointInstancerGrid/PointInstancer/%d' % instanceIndex)
])
ufeItem = ufe.Hierarchy.createItem(ufePath)
# Select the point instance scene item.
globalSelection = ufe.GlobalSelection.get()
globalSelection.append(ufeItem)
# Get the PointInstancer prim for validating the values in USD.
ufePathString = ufe.PathString.string(ufePath)
prim = mayaUsdUfe.ufePathToPrim(ufePathString)
pointInstancer = UsdGeom.PointInstancer(prim)
self.assertTrue(pointInstancer)
# The PointInstancer should not have any authored orientations
# initially.
orientationsAttr = pointInstancer.GetOrientationsAttr()
self.assertFalse(orientationsAttr.HasAuthoredValue())
# Perfom an orientation manipulation via the rotate command.
cmds.rotate(15.0, 30.0, 45.0, objectSpace=True, relative=True)
# The initial rotate should have filled out the orientations attribute.
orientations = orientationsAttr.Get()
self.assertEqual(len(orientations), 14)
# Validate the rotated item.
orientation = orientations[instanceIndex]
self.assertTrue(Gf.IsClose(orientation.real, 0.897461, self.EPSILON))
self.assertTrue(
Gf.IsClose(orientation.imaginary,
Gf.Vec3h(0.01828, 0.2854, 0.335205), self.EPSILON))
# The non-rotated items should all have identity orientations.
for i in [idx for idx in range(14) if idx != instanceIndex]:
orientation = orientations[i]
self.assertTrue(
Gf.IsClose(orientation.real,
Gf.Quath.GetIdentity().real, self.EPSILON))
self.assertTrue(
Gf.IsClose(orientation.imaginary,
Gf.Quath.GetIdentity().imaginary, self.EPSILON))
# Perfom another rotate.
cmds.rotate(25.0, 50.0, 75.0, objectSpace=True, relative=True)
# Re-fetch the USD orientation and check for the update.
orientation = orientationsAttr.Get()[instanceIndex]
self.assertTrue(Gf.IsClose(orientation.real, 0.397949, self.EPSILON))
self.assertTrue(
Gf.IsClose(orientation.imaginary,
Gf.Vec3h(-0.0886841, 0.57666, 0.708008), self.EPSILON))
# Now undo, and re-check.
cmds.undo()
orientation = orientationsAttr.Get()[instanceIndex]
self.assertTrue(Gf.IsClose(orientation.real, 0.897461, self.EPSILON))
self.assertTrue(
Gf.IsClose(orientation.imaginary,
Gf.Vec3h(0.01828, 0.2854, 0.335205), self.EPSILON))
# And once more.
# Note that with more complete undo support, this would ideally clear
# the authored orientations attribute, returning it to its true
# original state. For now, we just validate that the orientation value
# is back to its default of identity.
cmds.undo()
orientation = orientationsAttr.Get()[instanceIndex]
self.assertTrue(
Gf.IsClose(orientation.real,
Gf.Quath.GetIdentity().real, self.EPSILON))
self.assertTrue(
Gf.IsClose(orientation.imaginary,
Gf.Quath.GetIdentity().imaginary, self.EPSILON))
