Assign Avalon/Colorbleed-config published looks on VRayProxy nodes that load Alembic files.

maya_vrayproxy_assign_shaders_avalon.py

import os
from collections import defaultdict
import logging
import alembic.Abc
import json
from maya import cmds

import colorbleed.maya.lib as lib
import avalon.io as io
import avalon.maya
import avalon.api as api


log = logging.getLogger(__name__)


def get_alembic_paths_by_property(filename, attr, verbose=False):
    """Return attribute value per objects in the Alembic file.

    Reads an Alembic archive hierarchy and retrieves the
    value from the `attr` properties on the objects.

    Arguments:
        filename (str): Full path to Alembic archive to read.
        verbose (bool): Whether to verbosely log missing attributes.

    Returns:
        dict: Mapping of node full path with its id

    """

    # Normalize alembic path
    filename = os.path.normpath(filename)
    filename = filename.replace("\\", "/")
    filename = str(filename)  # path must be string

    archive = alembic.Abc.IArchive(filename)
    root = archive.getTop()

    iterator = list(root.children)
    obj_ids = dict()

    for obj in iterator:
        name = obj.getFullName()

        # include children for coming iterations
        iterator.extend(obj.children)

        props = obj.getProperties()
        if props.getNumProperties() == 0:
            # Skip those without properties, e.g. '/materials' in a gpuCache
            continue

        # THe custom attribute is under the properties' first container under
        # the ".arbGeomParams"        
        prop = props.getProperty(0)     # get base property

        property = None
        try:
            geomParams = prop.getProperty('.arbGeomParams')
            property = geomParams.getProperty(attr)
        except KeyError:
            if verbose:
                log.debug("Missing attr on: {0}".format(name))
            continue

        if not property.isConstant():
            log.warning("Id not constant on: {0}".format(name))
            
        # Get first value sample
        value = property.getValue()[0]

        obj_ids[name] = value

    return obj_ids


def get_alembic_ids_cache(path):
    """Build a id to node mapping in Alembic file

    Nodes without IDs are ignored.

    Returns:
        dict: Mapping of id to nodes in the Alembic.

    """

    node_ids = get_alembic_paths_by_property(path, attr="cbId")
    id_nodes = defaultdict(list)
    for node, id in node_ids.iteritems():
        id_nodes[id].append(node)

    return dict(id_nodes.iteritems())
    
    
def assign_vrayproxy_shaders(proxy, assignments):
    # todo: allow to optimize and assign a single shader to multiple shapes at once?
    #       or maybe even set it to the highest available path?
    
    # Clear all current shader assignments
    plug = proxy + ".shaders"
    num = cmds.getAttr(plug, size=True)
    for i in reversed(range(num)):
        cmds.removeMultiInstance("{}[{}]".format(plug, i), b=True)
    
    # Create new assignment overrides
    index = 0
    for material, paths in assignments.items():
        for path in paths:
            plug = "{}.shaders[{}]".format(proxy, index)
            cmds.setAttr(plug + ".shadersNames", path, type="string")
            cmds.connectAttr(material + ".outColor", 
                             plug + ".shadersConnections", force=True)
            index += 1
        

def get_look_relationships(version_id):

    json_representation = io.find_one({"type": "representation",
                                       "parent": version_id,
                                       "name": "json"})

    # Load relationships
    shader_relation = api.get_representation_path(json_representation)
    with open(shader_relation, "r") as f:
        relationships = json.load(f)

    return relationships
    
    
def load_look(version_id, use_existing=True):

    # Get representations of shader file and relationships
    look_representation = io.find_one({"type": "representation",
                                       "parent": version_id,
                                       "name": "ma"})

    # See if representation is already loaded, if so reuse it.
    host = api.registered_host()
    representation_id = str(look_representation['_id'])
    for container in host.ls():
        if (container['loader'] == "LookLoader" and
                container['representation'] == representation_id):
            log.info("Reusing loaded look ..")
            container_node = container['objectName']
            break
    else:
        log.info("Using look for the first time ..")

        # Load file
        loaders = api.loaders_from_representation(api.discover(api.Loader),
                                                  representation_id)
        Loader = next((i for i in loaders if i.__name__ == "LookLoader"), None)
        if Loader is None:
            raise RuntimeError("Could not find LookLoader, this is a bug")

        # Reference the look file
        with avalon.maya.maintained_selection():
            container_node = api.load(Loader, look_representation)

    # Get container members
    shader_nodes = cmds.sets(container_node, query=True)
    return shader_nodes


def get_latest_version(asset_id, subset):

    subset = io.find_one({"name": subset, 
                          "parent": io.ObjectId(asset_id), 
                           "type": "subset"})
    if not subset:
        raise RuntimeError("Subset does not exist: %s" % subset)
                 
    version = io.find_one({"type": "version",
                           "parent": subset["_id"]},
                           sort=[("name", -1)])
    if not version:
        raise RuntimeError("Version does not exist.")

    return version
    
    
def vrayproxy_assign_look(proxy, subset="lookDefault"):

    path = cmds.getAttr(proxy + ".fileName")

    nodes_by_id = get_alembic_ids_cache(path)
    if not nodes_by_id:
        log.warning("Alembic file has no cbId attributes: %s" % path)
        return
    
    # Group by asset id so we run over the look per asset
    node_ids_by_asset_id = defaultdict(set)
    for node_id in nodes_by_id:
        asset_id = node_id.split(":", 1)[0]
        node_ids_by_asset_id[asset_id].add(node_id)
        
    for asset_id, node_ids in node_ids_by_asset_id.items():
        
        # Get latest look version
        try:
            version = get_latest_version(asset_id, subset=subset)
        except RuntimeError as exc:
            print(exc)
            continue
        
        relationships = get_look_relationships(version["_id"])
        shadernodes = load_look(version["_id"])
            
        # Get only the node ids and paths related to this asset
        # And get the shader edits the look supplies
        asset_nodes_by_id = {node_id: nodes_by_id[node_id] for node_id in node_ids}
        edits = list(lib.iter_shader_edits(relationships, shadernodes, asset_nodes_by_id))
        
        # Create assignments
        assignments = {}
        for edit in edits:
            if edit["action"] == "assign":
                nodes = edit["nodes"]
                shader = edit["shader"]
                if not cmds.ls(shader, type="shadingEngine"):
                    print("Skipping non-shader: %s" % shader)
                    continue
                
                inputs = cmds.listConnections(shader + ".surfaceShader", source=True)
                if not inputs:
                    print("Shading engine missing material: %s" % shader)
                    
                # Strip off component assignments
                for i, node in enumerate(nodes):
                    if "." in node:
                        log.warning("Converting face assignment to full object assignment. This conversion can be lossy: %s" % node)
                        nodes[i] = node.split(".")[0]
                    
                material = inputs[0]
                assignments[material] = nodes
            
        assign_vrayproxy_shaders(proxy, assignments)


# Example usage
if __name__ == "__main__":

    # Ensure V-Ray is loaded
    cmds.loadPlugin("vrayformaya", quiet=True)
    
    # Assign lookDefault to all V-Ray Proxies
    for proxy in cmds.ls(sl=True, dag=True, type="VRayProxy"):
        vrayproxy_assign_look(proxy, subset="lookDefault")