PoseRenderer/python/openpose_renderer/draw_override.py

"""Viewport 2.0 draw override for ``openposeCharacter``: a live, colored
skeleton overlay drawn in world space on top of the mapped joints, so a rig
can be visually QA'd against the OpenPose_full mapping without rendering.

This is independent of the render override (render_override.py) -- it always
draws (when the node's enableDisplay attr is on) in every viewport, in 3D
world space, and tracks animation live via prepareForDraw() being re-run
whenever the node is dirtied.
"""

import maya.api.OpenMaya as om
import maya.api.OpenMayaRender as omr

from . import constants
from . import face_blendshapes
from . import mapping_node
from . import projector


def _mcolor(rgb255):
    return om.MColor((rgb255[0] / 255.0, rgb255[1] / 255.0, rgb255[2] / 255.0, 1.0))


class OpenPoseUserData(om.MUserData):

    def __init__(self):
        om.MUserData.__init__(self, False)  # Maya should not delete after one draw
        self.lines = []   # [(p0, p1, MColor, thickness), ...]
        self.points = []  # [(p, MColor, pointSize), ...]


class OpenPoseDrawOverride(omr.MPxDrawOverride):

    @staticmethod
    def creator(obj):
        return OpenPoseDrawOverride(obj)

    def __init__(self, obj):
        omr.MPxDrawOverride.__init__(self, obj, None, isAlwaysDirty=True)

    def supportedDrawAPIs(self):
        return omr.MRenderer.kAllDevices

    def isBounded(self, obj_path, camera_path):
        # Joints being displayed can be anywhere in the scene relative to
        # this node's own (irrelevant) transform -- don't cull on bbox.
        return False

    def hasUIDrawables(self):
        return True

    def prepareForDraw(self, obj_path, camera_path, frame_context, old_data):
        data = old_data if isinstance(old_data, OpenPoseUserData) else OpenPoseUserData()
        data.lines = []
        data.points = []

        dep_fn = om.MFnDependencyNode(obj_path.node())
        try:
            enabled = dep_fn.findPlug(mapping_node.aEnableDisplay, False).asBool()
        except RuntimeError:
            enabled = True
        if not enabled:
            return data

        radius_scale = dep_fn.findPlug(mapping_node.aJointRadiusScale, False).asFloat()
        thickness_scale = dep_fn.findPlug(mapping_node.aLimbThicknessScale, False).asFloat()
        body_model = mapping_node.get_body_model(dep_fn)

        mapped = mapping_node.get_mapped_joint_paths(dep_fn)
        positions = {index: projector.world_position(path) for index, path in mapped.items()}

        if mapping_node.get_face_source(dep_fn) == mapping_node.FACE_SOURCE_ARKIT_BLENDSHAPES:
            nose_index = constants.NAME_TO_INDEX["Body_Nose"]
            anchor = positions.get(nose_index)
            if anchor is not None:
                weights = mapping_node.get_blendshape_weights(dep_fn)
                face_scale = mapping_node.get_face_scale(dep_fn)
                local_points = face_blendshapes.compute_face_local_points(weights)
                world_points = projector.billboard_world_points(anchor, camera_path, local_points, face_scale)
                for offset, world_point in enumerate(world_points):
                    positions[constants.FACE_START + offset] = world_point

        active_indices = set(constants.active_point_indices(body_model))
        pairs, colors, thickness_list = constants.full_limb_data(body_model)

        for pair_i, (a, b) in enumerate(pairs):
            if a in positions and b in positions:
                thickness = thickness_list[pair_i] * thickness_scale
                if thickness <= 0.0:
                    continue
                data.lines.append((positions[a], positions[b], _mcolor(colors[pair_i]), thickness))

        for index, pos in positions.items():
            if index not in active_indices:
                continue
            radius = constants.KEYPOINT_POINT_RADIUS[index] * radius_scale
            if radius <= 0.0:
                continue
            data.points.append((pos, _mcolor(constants.KEYPOINT_POINT_COLORS[index]), radius * 2.0))

        return data

    def addUIDrawables(self, obj_path, draw_manager, frame_context, data):
        if not isinstance(data, OpenPoseUserData):
            return

        draw_manager.beginDrawable()
        draw_manager.setDepthPriority(omr.MRenderItem.sActiveWireDepthPriority)

        for p0, p1, color, thickness in data.lines:
            draw_manager.setColor(color)
            draw_manager.setLineWidth(thickness)
            draw_manager.line(p0, p1)

        for pos, color, point_size in data.points:
            draw_manager.setColor(color)
            draw_manager.setPointSize(point_size)
            draw_manager.point(pos)

        draw_manager.endDrawable()


def register(fn_plugin):
    omr.MDrawRegistry.registerDrawOverrideCreator(
        mapping_node.DRAW_CLASSIFICATION,
        mapping_node.DRAW_REGISTRANT_ID,
        OpenPoseDrawOverride.creator,
    )


def deregister(fn_plugin):
    omr.MDrawRegistry.deregisterDrawOverrideCreator(
        mapping_node.DRAW_CLASSIFICATION,
        mapping_node.DRAW_REGISTRANT_ID,
    )