PoseRenderer/python/openpose_renderer/mapping_node.py

"""``openposeCharacter`` node: per-character OpenPose_full joint mapping.

A pure MPxLocatorNode with one connectable message-attribute slot per
OpenPose_full keypoint (``targetJoint[0..136]``). The mapping lives entirely
as Maya connections (joint.message -> openposeCharacter.targetJoint[i]) so it
survives joint renames/reparenting and round-trips through the scene file
with no separate config to keep in sync. The node has no compute() output --
it is a passive connection target read directly by the draw override
(live viewport skeleton) and the render override (OpenPose image output).
"""

import maya.api.OpenMaya as om
import maya.api.OpenMayaUI as omui

from . import constants
from . import face_blendshapes

NODE_NAME = "openposeCharacter"
NODE_ID = om.MTypeId(0x0012F5C0)
DRAW_CLASSIFICATION = "drawdb/geometry/openposeCharacter"
DRAW_REGISTRANT_ID = "openposeRendererPlugin"

# Attribute handles, filled in by initialize()
aTargetJoint = None
aEnableDisplay = None
aJointRadiusScale = None
aLimbThicknessScale = None
aCharacterName = None
aBodyModel = None
aFaceSource = None
aBlendshapeWeight = None
aFaceScale = None
aFaceFollow = None
aUseDefaultEyesEars = None

# aBodyModel enum field values, in registration order (index 0 = BODY_25).
BODY_MODEL_FIELDS = [constants.BODY_MODEL_BODY_25, constants.BODY_MODEL_COCO]

# aFaceSource enum field values, in registration order (index 0 = Joints).
FACE_SOURCE_JOINTS = "Joints"
FACE_SOURCE_ARKIT_BLENDSHAPES = "ARKitBlendshapes"
FACE_SOURCE_FIELDS = [FACE_SOURCE_JOINTS, FACE_SOURCE_ARKIT_BLENDSHAPES]

# aFaceFollow enum field values, in registration order (index 0 = Camera).
FACE_FOLLOW_CAMERA = "Camera"
FACE_FOLLOW_NECK = "Neck"
FACE_FOLLOW_FIELDS = [FACE_FOLLOW_CAMERA, FACE_FOLLOW_NECK]

DEFAULT_FACE_SCALE = 15.0


class OpenposeCharacterNode(omui.MPxLocatorNode):

    kTypeName = NODE_NAME
    kTypeId = NODE_ID

    def __init__(self):
        omui.MPxLocatorNode.__init__(self)

    def compute(self, plug, data_block):
        # Pure passive mapping node: nothing to compute.
        return None

    def isBounded(self):
        return False

    @staticmethod
    def creator():
        return OpenposeCharacterNode()

    @staticmethod
    def initialize():
        global aTargetJoint, aEnableDisplay, aJointRadiusScale
        global aLimbThicknessScale, aCharacterName, aBodyModel
        global aFaceSource, aBlendshapeWeight, aFaceScale, aFaceFollow
        global aUseDefaultEyesEars

        msg_attr = om.MFnMessageAttribute()
        aTargetJoint = msg_attr.create("targetJoint", "tj")
        msg_attr.array = True
        # indexMatters already defaults to True (sparse, index-preserving
        # array) -- explicitly assigning True here throws kInvalidParameter
        # in Maya's Python binding, so it is deliberately left untouched.
        msg_attr.storable = True
        msg_attr.readable = True
        msg_attr.writable = True
        om.MPxNode.addAttribute(aTargetJoint)

        num_attr = om.MFnNumericAttribute()

        aEnableDisplay = num_attr.create(
            "enableDisplay", "ed", om.MFnNumericData.kBoolean, True
        )
        num_attr.keyable = True
        om.MPxNode.addAttribute(aEnableDisplay)

        aJointRadiusScale = num_attr.create(
            "jointRadiusScale", "jrs", om.MFnNumericData.kFloat, 1.0
        )
        num_attr.keyable = True
        num_attr.setMin(0.0)
        om.MPxNode.addAttribute(aJointRadiusScale)

        aLimbThicknessScale = num_attr.create(
            "limbThicknessScale", "lts", om.MFnNumericData.kFloat, 1.0
        )
        num_attr.keyable = True
        num_attr.setMin(0.0)
        om.MPxNode.addAttribute(aLimbThicknessScale)

        str_data = om.MFnStringData()
        empty_string_obj = str_data.create("")
        typed_attr = om.MFnTypedAttribute()
        aCharacterName = typed_attr.create(
            "characterName", "cn", om.MFnData.kString, empty_string_obj
        )
        typed_attr.storable = True
        om.MPxNode.addAttribute(aCharacterName)

        enum_attr = om.MFnEnumAttribute()
        aBodyModel = enum_attr.create("bodyModel", "bmd", 0)
        for field_index, field_name in enumerate(BODY_MODEL_FIELDS):
            enum_attr.addField(field_name, field_index)
        enum_attr.keyable = True
        om.MPxNode.addAttribute(aBodyModel)

        face_source_attr = om.MFnEnumAttribute()
        aFaceSource = face_source_attr.create("faceSource", "fsrc", 0)
        for field_index, field_name in enumerate(FACE_SOURCE_FIELDS):
            face_source_attr.addField(field_name, field_index)
        face_source_attr.keyable = True
        om.MPxNode.addAttribute(aFaceSource)

        weight_attr = om.MFnNumericAttribute()
        aBlendshapeWeight = weight_attr.create("blendshapeWeight", "bsw", om.MFnNumericData.kFloat, 0.0)
        weight_attr.array = True
        weight_attr.keyable = True
        weight_attr.setMin(0.0)
        weight_attr.setMax(1.0)
        om.MPxNode.addAttribute(aBlendshapeWeight)

        face_scale_attr = om.MFnNumericAttribute()
        aFaceScale = face_scale_attr.create("faceScale", "fsc", om.MFnNumericData.kFloat, DEFAULT_FACE_SCALE)
        face_scale_attr.keyable = True
        face_scale_attr.setMin(0.0)
        om.MPxNode.addAttribute(aFaceScale)

        face_follow_attr = om.MFnEnumAttribute()
        aFaceFollow = face_follow_attr.create("faceFollow", "ffol", 0)
        for field_index, field_name in enumerate(FACE_FOLLOW_FIELDS):
            face_follow_attr.addField(field_name, field_index)
        face_follow_attr.keyable = True
        om.MPxNode.addAttribute(aFaceFollow)

        default_eyes_ears_attr = om.MFnNumericAttribute()
        aUseDefaultEyesEars = default_eyes_ears_attr.create(
            "useDefaultEyesEars", "udee", om.MFnNumericData.kBoolean, True
        )
        default_eyes_ears_attr.keyable = True
        om.MPxNode.addAttribute(aUseDefaultEyesEars)


def get_mapped_joint_paths(dep_node_fn):
    """Return {keypoint_index: MDagPath} for every connected targetJoint slot.

    ``dep_node_fn`` is an MFnDependencyNode wrapping an openposeCharacter.
    Unmapped keypoint indices are simply absent from the returned dict.
    """
    result = {}
    plug = dep_node_fn.findPlug(aTargetJoint, False)
    for logical_index in plug.getExistingArrayAttributeIndices():
        if logical_index < 0 or logical_index >= constants.TOTAL_KEYPOINTS:
            continue
        element_plug = plug.elementByLogicalIndex(logical_index)
        if not element_plug.isConnected:
            continue
        source_plug = element_plug.source()
        if source_plug.isNull:
            continue
        source_node = source_plug.node()
        if not source_node.hasFn(om.MFn.kDagNode):
            continue
        dag_fn = om.MFnDagNode(source_node)
        result[logical_index] = dag_fn.getPath()
    return result


def get_body_model(dep_node_fn):
    """Return constants.BODY_MODEL_BODY_25 / BODY_MODEL_COCO for this character."""
    try:
        field_index = dep_node_fn.findPlug(aBodyModel, False).asShort()
    except RuntimeError:
        field_index = 0
    if 0 <= field_index < len(BODY_MODEL_FIELDS):
        return BODY_MODEL_FIELDS[field_index]
    return constants.BODY_MODEL_BODY_25


def get_face_source(dep_node_fn):
    """Return FACE_SOURCE_JOINTS / FACE_SOURCE_ARKIT_BLENDSHAPES for this character."""
    try:
        field_index = dep_node_fn.findPlug(aFaceSource, False).asShort()
    except RuntimeError:
        field_index = 0
    if 0 <= field_index < len(FACE_SOURCE_FIELDS):
        return FACE_SOURCE_FIELDS[field_index]
    return FACE_SOURCE_JOINTS


def get_blendshape_weights(dep_node_fn):
    """Return {arkit_blendshape_name: weight} for all 52 slots (0.0 default)."""
    plug = dep_node_fn.findPlug(aBlendshapeWeight, False)
    weights = {}
    for index, name in enumerate(face_blendshapes.ARKIT_BLENDSHAPE_NAMES):
        element_plug = plug.elementByLogicalIndex(index)
        weights[name] = element_plug.asFloat()
    return weights


def get_face_scale(dep_node_fn):
    try:
        return dep_node_fn.findPlug(aFaceScale, False).asFloat()
    except RuntimeError:
        return DEFAULT_FACE_SCALE


def get_face_follow(dep_node_fn):
    """Return FACE_FOLLOW_CAMERA / FACE_FOLLOW_NECK for this character."""
    try:
        field_index = dep_node_fn.findPlug(aFaceFollow, False).asShort()
    except RuntimeError:
        field_index = 0
    if 0 <= field_index < len(FACE_FOLLOW_FIELDS):
        return FACE_FOLLOW_FIELDS[field_index]
    return FACE_FOLLOW_CAMERA


def get_use_default_eyes_ears(dep_node_fn):
    try:
        return dep_node_fn.findPlug(aUseDefaultEyesEars, False).asBool()
    except RuntimeError:
        return True


def iter_character_nodes():
    """Yield an MFnDependencyNode for every openposeCharacter node in the scene."""
    it = om.MItDependencyNodes(om.MFn.kPluginLocatorNode)
    while not it.isDone():
        obj = it.thisNode()
        fn = om.MFnDependencyNode(obj)
        if fn.typeName == NODE_NAME:
            yield fn
        it.next()


def register(fn_plugin):
    fn_plugin.registerNode(
        NODE_NAME,
        NODE_ID,
        OpenposeCharacterNode.creator,
        OpenposeCharacterNode.initialize,
        om.MPxNode.kLocatorNode,
        DRAW_CLASSIFICATION,
    )


def deregister(fn_plugin):
    fn_plugin.deregisterNode(NODE_ID)