USD Stage Traversal

A howto guide for finding the right prims in an OpenUSD stage, using python

USD Stage Traversal

If you know the path to a prim you want, getting it is as easy as stage.GetPrimAtPath('/path/to/prim'). It's often the case that you want to find prims that match some condition, regardless of path. In that case you can traverse the stage searching for the interesting ones.

Setting the Stage

There are several apis for traversing stages. Let's make a simple usd file we can use for examples. I'm not going to populate this with data, just make a skeleton with all default values. Nothing will draw but we can still see how to walk over it.

Running this python code

from pxr import Kind, Usd, UsdGeom, UsdShade

stage = Usd.Stage.CreateInMemory()
scene_xform = UsdGeom.Xform.Define(stage, '/Scene')
scene_xform.GetPrim().SetKind(Kind.Tokens.group)

model_one = UsdGeom.Xform.Define(stage, '/Scene/ModelOne')
model_one.GetPrim().SetKind(Kind.Tokens.component)
UsdGeom.Mesh.Define(stage, '/Scene/ModelOne/geo1')
UsdGeom.Mesh.Define(stage, '/Scene/ModelOne/geo2')

model_two = UsdGeom.Xform.Define(stage, '/Scene/ModelTwo')
model_two.GetPrim().SetKind(Kind.Tokens.component)
UsdGeom.Mesh.Define(stage, '/Scene/ModelTwo/geo1')
UsdGeom.Mesh.Define(stage, '/Scene/ModelTwo/geo2')

mat = UsdShade.Material.Define(stage, '/Scene/Material')
binding_api = UsdShade.MaterialBindingAPI.Apply(model_one.GetPrim())
binding_api.Bind(mat)

stage.Export('./test-layer.usda')

Generates the file below.

#usda 1.0

def Xform "Scene" (
    kind = "group"
)
{
    def Xform "ModelOne" (
        apiSchemas = ["MaterialBindingAPI"]
        kind = "component"
    )
    {
        rel material:binding = </Scene/Material>

        def Mesh "geo1"
        {
        }

        def Mesh "geo2"
        {
        }
    }

    def Xform "ModelTwo" (
        kind = "component"
    )
    {
        def Mesh "geo1"
        {
        }

        def Mesh "geo2"
        {
        }
    }

    def Material "Material"
    {
    }
}

The layer we generated exported to usda

I won't go too in depth on the python part. It's just there in case you want to recreate this file or make different versions for experimenting. I do want to point out the kind metadata on several prims. That's for the Model traversals we'll cover later on.

Walk Over All Prims

If we just want to see everything that's in the stage, the stage object has a few functions that return iterators. stage.Traverse walks over all prims in depth first order.

>>> for prim in stage.Traverse():
...     print('{0: <10}'.format(prim.GetTypeName()), prim.GetPath())
... 
Xform      /Scene
Xform      /Scene/ModelOne
Mesh       /Scene/ModelOne/geo1
Mesh       /Scene/ModelOne/geo2
Xform      /Scene/ModelTwo
Mesh       /Scene/ModelTwo/geo1
Mesh       /Scene/ModelTwo/geo2
Material   /Scene/Material

stage.TraverseAll() is also an option. In our example it has the same output as Traverse. Traverse excludes some things by default:

  • inactive prims: Prims with active=false behave like they've been deleted, so they don't show up in Traverse
  • unloaded prims: Payloads that haven't been loaded aren't in the scene so Traverse won't show their contents
  • undefined prims: If a prim has an over but no def then it is "undefined". It doesn't affect the scene, so Traverse won't show it
  • abstract prims: A prim is abstract if it is a class. Classes are types, not content, so they don't show up in Traverse

TraverseAll would show those things.

stage.Traverse can also take an optional parameter to limit the traversal. This parameter is called a Usd_PrimFlagsPredicate. I always find the documentation for these hard to find, so here's a link to the header where they're defined.

If we wanted to only look at prims tagged as Models using kinds, we could pass the Usd.PrimIsModel predicate like this

>>> for prim in stage.Traverse(Usd.PrimIsModel):
...     print('{0: <10}'.format(prim.GetTypeName()), prim.GetPath())
... 
Xform      /Scene
Xform      /Scene/ModelOne
Xform      /Scene/ModelTwo

Only these prims have Model kinds in our stage. The Scene prim is a model group, which is considered a model, and the two Model prims are component models in the Scene group.

We can also get only groups using Usd.PrimIsGroup

>>> for prim in stage.Traverse(Usd.PrimIsGroup):
...     print('{0: <10}'.format(prim.GetTypeName()), prim.GetPath())
...
Xform      /Scene

stage.Traverse(Usd.PrimAllPrimsPredicate) is equivalent to stage.TraverseAll()

Usd Prim Range

If you're following along in python you may have noticed that the object stage.Traverse returns is of type Usd.PrimRange. UsdPrimRange objects can also be created directly.

You can recreate the stage.Traverse behavior with a UsdPrimRange directly like

  • stage.Traverse() => Usd.PrimRange.Stage(stage)
  • stage.Traverse(Usd.PrimIsModel) => Usd.PrimRange.Stage(stage, Usd.PrimIsModel)
  • stage.TraverseAll() => Usd.PrimRange.Stage(stage, Usd.PrimAllPrimsPredicate)

Prim ranges also offer other constructors for more specific use cases. For instance, you can traverse starting at a specific prim instead of the whole stage. This code will limit the traversal to the contents of /Scene/ModelOne

>>> model_one = stage.GetPrimAtPath('/Scene/ModelOne')
>>> r = Usd.PrimRange(model_one)
>>> for prim in r:
...     print('{0: <10}'.format(prim.GetTypeName()), prim.GetPath())
... 
Xform      /Scene/ModelOne
Mesh       /Scene/ModelOne/geo1
Mesh       /Scene/ModelOne/geo2

With a prim range you can prune subtrees, basically stop traversing at a prim when you know none of the children are interesting. To do this you need to wrap the UsdPrimRange in an iter. The code below will traverse the stage but skip component model contents.

>>> r = iter(Usd.PrimRange.Stage(stage))
>>> for prim in r:
...     if Usd.ModelAPI(prim).GetKind() == Kind.Tokens.component:
...         r.PruneChildren()
...     print('{0: <10}'.format(prim.GetTypeName()), prim.GetPath())
... 
Xform      /Scene
Xform      /Scene/ModelOne
Xform      /Scene/ModelTwo
Material   /Scene/Material

UsdPrimRange has some interesting constructors that allow doing "pre and post" style traversals of the tree. This also requires using an iter wrapper. If we wanted to print the prim names nested beneath their parents, we could do so like this

>>> r = iter(Usd.PrimRange.PreAndPostVisit(stage.GetPrimAtPath('/Scene')))
>>> indent = 0
>>> for prim in r:
...     if r.IsPostVisit():
...         indent = indent - 1
...     else:
...         print("  "*indent, prim.GetName())
...         indent = indent + 1
... 
 Scene
   ModelOne
     geo1
     geo2
   ModelTwo
     geo1
     geo2
   Material

Just for completeness' sake, the prim function prim.GetFilteredChildren(predicate) can do a lot of the same things a prim range can, in case it is ever convenient.

Recursive Traversal

Truth be told, I often just traverse the stage by recursively calling GetChildren. The PreAndPostVisit example above could also be recreated like this.

>>> def visit_children(prim, indent):
...     print(indent, prim.GetName())
...     for child in prim.GetChildren():
...         visit_children(child, indent + '  ')
... 
>>> visit_children(stage.GetPrimAtPath('/Scene'), '')
 Scene
   ModelOne
     geo1
     geo2
   ModelTwo
     geo1
     geo2
   Material

Materials Example

Let's gather some information about the Material prim. It's directly bound to ModelOne, so it will be the material for the two meshes in that model, but they don't have direct bindings. Let's answer two questions about our stage.

A) What prims are bound to Material?
B) What meshes are shaded with Material?

This code will walk the stage and answer these questions.

>>> for prim in stage.Traverse():
...     materialAPI = UsdShade.MaterialBindingAPI(prim)
...     direct = materialAPI.GetDirectBinding().GetMaterialPath()
...     if direct:
...         print(prim.GetPath(), "has a direct binding to", direct)
...     computed = materialAPI.ComputeBoundMaterial()
...     if UsdGeom.Mesh(prim) and computed and len(computed) > 0 and computed[0]:
...         print(prim.GetPath(), "uses the material", computed[0].GetPath())
... 
/Scene/ModelOne has a direct binding to /Scene/Material
/Scene/ModelOne/geo1 uses the material /Scene/Material
/Scene/ModelOne/geo2 uses the material /Scene/Material

This is a pretty simple example and it ignores some details, it wouldn't work for collection based binding, doesn't handle multiple materials, only looks for materials on meshes, and I'm sure lots of other issues. A general function for this would be more complicated, but I hope this demonstrates the ways you might use the traversal code in scripts.