Skip to content

Instantly share code, notes, and snippets.

@BigRoy
Created October 19, 2023 13:45
Show Gist options
  • Save BigRoy/df6e6a2bb1538fe0589a330582bad54e to your computer and use it in GitHub Desktop.
Save BigRoy/df6e6a2bb1538fe0589a330582bad54e to your computer and use it in GitHub Desktop.
Maya USD Import Chaser in Python to import `primvars:cbId_transform` attribute to the transform as `cbId` instead of as attribute on the shape
import logging
import mayaUsd.lib as mayaUsdLib
from maya import cmds
import maya.api.OpenMaya as OpenMaya
def create_cbid_attr(default_value):
default = OpenMaya.MFnStringData().create(default_value)
return OpenMaya.MFnTypedAttribute().create("cbId",
"cbId",
OpenMaya.MFnData.kString,
default)
class ImportChaserCbIdShapes(mayaUsdLib.ImportChaser):
"""Import `cbId_transform` attribute to `cbId` on transform node"""
def __init__(self, factoryContext, *args, **kwargs):
super(ImportChaserCbIdShapes, self).__init__(factoryContext,
*args,
**kwargs)
self.modifier = OpenMaya.MDGModifier()
def _post_import(self, returnPredicate, stage, dagPath, sdfPath, jobArgs):
fn = OpenMaya.MFnDependencyNode()
# Note: The Sdf to Dag map only contains the dag entries for which
# an SdfPath exists. Since shapes are expanded from its parent
# they do not get their own mapping. As such, we need to manually
# detect whether we're dealing with a shape currently
for mapping in self.GetSdfToDagMap():
dag_path = mapping.data()
sdf_path = mapping.key()
prim = stage.GetPrimAtPath(sdf_path)
if not prim.HasAttribute("primvars:cbId_transform"):
continue
num = dag_path.numberOfShapesDirectlyBelow()
if num == 0:
continue
if num > 1:
logger.warning("More than one shape for: {}".format(dag_path))
# Get the shape
shape_dag_path = OpenMaya.MDagPath(dag_path).extendToShape(0)
shape_node = shape_dag_path.node()
fn.setObject(shape_node)
# Remove cbId_transform from shape if it exists tehre
plug = fn.hasAttribute("cbId_transform")
if plug:
shape_attr = fn.attribute("cbId_transform")
self.modifier.removeAttribute(shape_node, shape_attr)
# Add to transform instead
value = prim.GetAttribute("primvars:cbId_transform").Get() or ""
attr = create_cbid_attr(value)
transform_node = dag_path.node()
self.modifier.addAttribute(transform_node, attr)
self.modifier.doIt()
return True
def PostImport(self, returnPredicate, stage, dagPaths, sdfPaths, jobArgs):
try:
return self._post_import(
returnPredicate,
stage,
dagPaths,
sdfPaths,
jobArgs
)
except Exception as exc:
logging.error(exc, exc_info=True)
raise
def Redo(self):
# NOTE: (yliangsiew) Undo the undo to re-do.
self.modifier.doIt()
return True
def Undo(self):
self.modifier.undoIt()
return True
mayaUsdLib.ImportChaser.Register(ImportChaserCbIdShapes, "cbId_transform")
path = r"P:\Projects\test\production\usd_cbId_test\work\FX\houdini\geo\with_suffix_roy.usda"
nodes = cmds.file(path,
sharedReferenceFile=False,
reference=True,
returnNewNodes=True,
options="chaser=[cbId_transform]")
#rootPaths = cmds.mayaUSDImport(v=True,
# f=path,
# chaser=['cbId_transform'])
@BigRoy
Copy link
Author

BigRoy commented Oct 23, 2023

Maya USD export chaser to write primvars:cbId_transform and primvars:cbId attributes

import logging

from maya.api import OpenMaya
import mayaUsd.lib as mayaUsdLib
from pxr import Sdf


class CbExportChaser(mayaUsdLib.ExportChaser):

    def __init__(self, factoryContext, *args, **kwargs):
        super(CbExportChaser, self).__init__(factoryContext, *args, **kwargs)
        self.dag_to_usd = factoryContext.GetDagToUsdMap()
        self.stage = factoryContext.GetStage()
        self.job_args = factoryContext.GetJobArgs()

    def PostExport(self):
        try:
            return self._post_export()
        except Exception as exc:
            logging.error(exc, exc_info=True)
            raise

    def _post_export(self):
        
        merged_shapes = getattr(self.job_args, "mergeTransformAndShape", True)
        
        fn = OpenMaya.MFnDependencyNode()
        for entry in self.dag_to_usd:
            prim_path = entry.data()
            dag_path = entry.key()

            fn.setObject(dag_path.node())
            if not fn.hasAttribute("cbId"):
                continue

            try:
                value = fn.findPlug("cbId", True).asString()
            except RuntimeError:
                logging.error(
                    "Unable to get {dag_path}.cbId value.".format(
                        dag_path=dag_path
                    ), 
                    exc_info=True
                )
                continue

            # Create attributes
            is_parent_of_shape = merged_shapes and bool(dag_path.numberOfShapesDirectlyBelow())
            attr_name = "cbId_transform" if is_parent_of_shape else "cbId" 
            attr_name = "primvars:{}".format(attr_name)

            prim = self.stage.GetPrimAtPath(prim_path)
            attr = prim.CreateAttribute(attr_name, Sdf.ValueTypeNames.String)
            attr.Set(value)

        return True


# Testing
from maya import cmds

mayaUsdLib.ExportChaser.Unregister(CbExportChaser, "cbid")
mayaUsdLib.ExportChaser.Register(CbExportChaser, "cbid")

cmds.loadPlugin("mayaUsdPlugin", quiet=True)
cmds.mayaUSDExport(file="E:/test.usda",
                   frameRange=(1, 5),
                   frameStride=1.0,
                   chaser=["cbid"],
                   # mergeTransformAndShape=True,
                   selection=True)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment