Maya: Convert AiStandardSurface to USDPreview Surface.

mcatu.py

"""
Will convert aiStandardSurface to usdPreviewSurface.

Usage:
- Select aiStandardSurface nodes and execute code.
"""
import pymel.core as pc


def list_connections(node, skip_nodes, **kwargs):
    connections = node.listConnections(**kwargs)

    results = []
    for dst, src in connections:
        if src.node().type() in skip_nodes:
            src = src.node().listConnections(**kwargs)[0][1]
        results.append([dst, src])

    return results

selection = pc.ls(selection=True)
attribute_mapping = {
    "baseColor": "diffuseColor",
    "normalCamera": "normal",
    "specularRoughness": "roughness",
    "specular": "specularColor",
    "specularColor": "specularColor",
    "transmission": "opacity",
    "metalness": "metallic"
}
skip_nodes = ["bump2d", "displacementShader"]
flagged_connections = []
for node in selection:
    # Create USD shader.
    usd_name = "USD_" + node.name()
    usd_mat = pc.shadingNode("usdPreviewSurface", asShader=True, name=usd_name)
    shading_engine = pc.sets(name="{}SG".format(usd_name), empty=True, renderable=True, noSurfaceShader=True)
    usd_mat.outColor >> shading_engine.surfaceShader

    # Basic connections.
    connections = list_connections(
        node, skip_nodes, destination=False, connections=True, source=True, plugs=True
    )
    connected_attributes = []
    for dest, source in connections:
        usd_attr_name = attribute_mapping[dest.attrName(longName=True)]
        usd_attr = pc.PyNode(
            "{}.{}".format(usd_mat.name(), usd_attr_name)
        )

        if source.type() != "file":
            flagged_connections.append([source, usd_attr])

        print("Connecting \"{}\" to \"{}\"".format(source, usd_attr))
        connected_attributes.append(usd_attr.attrName(longName=True))
        
        source_type = pc.getAttr(source, type=True)
        dest_type = pc.getAttr(usd_attr, type=True)
        if dest_type == "float3" and source_type == "float":
            for child_attr in usd_attr.children():
                source >> child_attr
            continue

        source >> usd_attr

    # Attribute transfer
    if node.metalness.get() == 0:
        usd_mat.useSpecularWorkflow.set(True)

    for k, v in attribute_mapping.items():
        if v in connected_attributes:
            continue

        source_name = "{}.{}".format(node.name(), k)
        source_attr = pc.PyNode(source_name)
        dest_name = "{}.{}".format(usd_mat.name(), attribute_mapping[k])
        dest_attr = pc.PyNode(dest_name)
        
        print("Setting \"{}\" > \"{}\"".format(source_attr, dest_attr))
        
        source_type = pc.getAttr(source_attr, type=True)
        dest_type = pc.getAttr(dest_attr, type=True)
        if source_type == dest_type:
            value = source_attr.get()
            
            # Clamp source value to dest range.
            min_value = dest_attr.getMin() or value
            max_value = dest_attr.getMax() or value
            value = max(min(value, max_value), min_value)
            
            # Transmission > Opacity is inverted.
            if k == "transmission":
                value = 1.0 - value

            dest_attr.set(value)
        elif dest_type == "float3" and source_type == "float":
            dest_attr.set(
                (source_attr.get(), source_attr.get(), source_attr.get())
            )
        else:
            raise NotImplementedError(
                "Cannot set \"{}\" ({}) from \"{}\" ({})".format(
                    dest_attr, dest_type, source_attr, source_type
                )
            )

    # Displacement
    shading_engine = node.listConnections(type="shadingEngine")[0]
    connections = shading_engine.displacementShader.listConnections(
        destination=False, connections=True, source=True, plugs=True
    )
    if connections:
        source = connections[0][1]
        if connections[0][1].node().type() == "displacementShader":        
            source = connections[0][1].node().listConnections(
                type="file", destination=False, connections=True, source=True, plugs=True
            )[0][1]
        source >> usd_mat.displacement
    
    # Assign material
    for assignment in pc.sets(shading_engine, q=True):
        pc.select(assignment)
        pc.hyperShade(a=usd_mat)

if flagged_connections:
    for src, dst in flagged_connections:
        print("Non-file connection between {} and {}".format(src, dst))
    cmds.confirmDialog(
        title="Conversion Error",
        message="Non-file nodes found. See script editor.",
        button=["Ok"],
        defaultButton="Ok"
    )