PoseRenderer/python/openpose_renderer/render_override.py

"""Maya renderer integration for OpenPose_full output.

Two pieces:

* ``OpenposeRenderOverride`` -- a Viewport 2.0 ``MRenderOverride`` registered
  as "OpenPose" in the viewport panel's Renderer dropdown, using the same
  minimal scene/HUD/present operation chain as Maya's own reference render
  override sample. This makes "OpenPose" a normal, selectable renderer
  rather than gambling on undocumented GPU render-target pixel readback
  from Python for the actual image content (see the plan's flagged risk).

* ``render_frame`` / the ``openposeRenderSequence`` command -- the actual
  OpenPose image generation: projects every ``openposeCharacter`` in the
  scene through a given camera and rasterizes them (projector.py +
  rasterizer.py, pure software, no GPU dependency) to a PNG per frame, and
  pushes each frame into Maya's Render View via
  ``renderWindowEditor(loadImage=...)`` -- the same mechanism third-party
  renderers use to display progress there, so the images are visible in the
  normal render workflow, not just written silently to disk.
"""

import os

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

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

UI_NAME = "OpenPose"
RENDER_OVERRIDE_NAME = "openposeRenderOverride"

_override_instance = None


class OpenposeRenderOverride(omr.MRenderOverride):

    def __init__(self, name):
        omr.MRenderOverride.__init__(self, name)
        self._operations = [
            omr.MSceneRender("openposeSceneRender"),
            omr.MHUDRender(),
            omr.MPresentTarget("openposePresentTarget"),
        ]
        self._current_operation = -1

    def uiName(self):
        return UI_NAME

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

    def setup(self, destination):
        return

    def cleanup(self):
        self._current_operation = -1

    def startOperationIterator(self):
        self._current_operation = 0
        return True

    def renderOperation(self):
        if 0 <= self._current_operation < len(self._operations):
            return self._operations[self._current_operation]
        return None

    def nextRenderOperation(self):
        self._current_operation += 1
        return self._current_operation < len(self._operations)


def _fill_default_eye_ear_keypoints(dep_fn, mapped, keypoints, camera_path, width, height):
    """Fill in synthetic REye/LEye/REar/LEar pixel keypoints (see
    constants.DEFAULT_EYE_EAR_LOCAL_OFFSETS) for whichever of those slots
    aren't already explicitly mapped, anchored at and oriented by the
    joint mapped to Nose (the head) -- so they turn with the head's actual
    rotation rather than facing the camera. No-op if Nose isn't mapped.
    """
    missing_offsets = {
        index: offset
        for index, offset in constants.DEFAULT_EYE_EAR_LOCAL_OFFSETS.items()
        if index not in mapped
    }
    if not missing_offsets:
        return
    nose_index = constants.NAME_TO_INDEX["Body_Nose"]
    if nose_index not in mapped:
        return
    anchor = projector.world_position(mapped[nose_index])

    face_scale = mapping_node.get_face_scale(dep_fn)
    computed = projector.default_head_relative_points(anchor, mapped[nose_index], face_scale, missing_offsets)
    for index, world_point in computed.items():
        px, py, in_front = projector.project_point(world_point, camera_path, width, height)
        keypoints[index] = (px, py, 1.0 if in_front else 0.0)


def render_frame(camera_path, width, height, output_path):
    """Project + rasterize every mapped openposeCharacter to a single PNG."""
    characters = []
    for dep_fn in mapping_node.iter_character_nodes():
        try:
            enabled = dep_fn.findPlug(mapping_node.aEnableDisplay, False).asBool()
        except RuntimeError:
            enabled = True
        if not enabled:
            continue
        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)
        keypoints = projector.project_character(mapped, camera_path, width, height)

        if mapping_node.get_use_default_eyes_ears(dep_fn):
            _fill_default_eye_ear_keypoints(dep_fn, mapped, keypoints, camera_path, width, height)

        if mapping_node.get_face_source(dep_fn) == mapping_node.FACE_SOURCE_ARKIT_BLENDSHAPES:
            nose_index = constants.NAME_TO_INDEX["Body_Nose"]
            nose_x, nose_y, nose_confidence = keypoints[nose_index]
            if nose_confidence > 0.0 and nose_index in mapped:
                anchor = projector.world_position(mapped[nose_index])
                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)

                neck_index = constants.NAME_TO_INDEX["Body_Neck"]
                if mapping_node.get_face_follow(dep_fn) == mapping_node.FACE_FOLLOW_NECK and neck_index in mapped:
                    world_points = projector.oriented_world_points(
                        anchor, mapped[neck_index], local_points, face_scale
                    )
                else:
                    world_points = projector.billboard_world_points(anchor, camera_path, local_points, face_scale)

                for offset, world_point in enumerate(world_points):
                    px, py, in_front = projector.project_point(world_point, camera_path, width, height)
                    keypoints[constants.FACE_START + offset] = (px, py, 1.0 if in_front else 0.0)

        characters.append((keypoints, radius_scale, thickness_scale, body_model))

    rasterizer.render_characters_to_file(characters, width, height, output_path)
    return output_path


def _default_output_dir():
    images_rule = cmds.workspace(fileRuleEntry="images", query=True) or "images"
    workspace_root = cmds.workspace(query=True, rootDirectory=True)
    out_dir = os.path.join(workspace_root, images_rule, "openpose")
    if not os.path.isdir(out_dir):
        os.makedirs(out_dir)
    return out_dir


def _resolve_camera_path(camera_name):
    if camera_name:
        selection = om.MSelectionList()
        selection.add(camera_name)
        dag_path = selection.getDagPath(0)
        if dag_path.apiType() == om.MFn.kTransform:
            dag_path.extendToShape()
        return dag_path
    import maya.api.OpenMayaUI as omui
    return omui.M3dView.active3dView().getCamera()


class RenderSequenceCommand(om.MPxCommand):

    kCmdName = "openposeRenderSequence"

    @staticmethod
    def creator():
        return RenderSequenceCommand()

    @staticmethod
    def createSyntax():
        syntax = om.MSyntax()
        syntax.addFlag("-s", "-startFrame", om.MSyntax.kDouble)
        syntax.addFlag("-e", "-endFrame", om.MSyntax.kDouble)
        syntax.addFlag("-cam", "-camera", om.MSyntax.kString)
        syntax.addFlag("-dir", "-outputDirectory", om.MSyntax.kString)
        syntax.addFlag("-sw", "-showInRenderView", om.MSyntax.kBoolean)
        return syntax

    def isUndoable(self):
        return False

    def doIt(self, args):
        parser = om.MArgParser(self.syntax(), args)

        start = (
            parser.flagArgumentDouble("-s", 0)
            if parser.isFlagSet("-s")
            else cmds.playbackOptions(query=True, minTime=True)
        )
        end = (
            parser.flagArgumentDouble("-e", 0)
            if parser.isFlagSet("-e")
            else cmds.playbackOptions(query=True, maxTime=True)
        )
        camera_name = parser.flagArgumentString("-cam", 0) if parser.isFlagSet("-cam") else None
        out_dir = parser.flagArgumentString("-dir", 0) if parser.isFlagSet("-dir") else _default_output_dir()
        show_in_render_view = parser.flagArgumentBool("-sw", 0) if parser.isFlagSet("-sw") else True

        if not os.path.isdir(out_dir):
            os.makedirs(out_dir)

        camera_path = _resolve_camera_path(camera_name)
        width, height = projector.get_render_resolution()

        scene_path = cmds.file(query=True, sceneName=True)
        scene_name = os.path.splitext(os.path.basename(scene_path))[0] if scene_path else "openpose"
        scene_name = scene_name or "openpose"

        if show_in_render_view:
            cmds.RenderViewWindow()

        original_time = cmds.currentTime(query=True)
        written = []
        frame = int(start)
        while frame <= int(end):
            cmds.currentTime(frame)
            file_path = os.path.join(out_dir, "{0}.{1:04d}.png".format(scene_name, frame))
            render_frame(camera_path, width, height, file_path)
            written.append(file_path)
            if show_in_render_view:
                cmds.renderWindowEditor("renderView", edit=True, loadImage=file_path)
            frame += 1
        cmds.currentTime(original_time)

        self.setResult(written)


def register(fn_plugin):
    global _override_instance
    fn_plugin.registerCommand(
        RenderSequenceCommand.kCmdName, RenderSequenceCommand.creator, RenderSequenceCommand.createSyntax
    )
    _override_instance = OpenposeRenderOverride(RENDER_OVERRIDE_NAME)
    omr.MRenderer.registerOverride(_override_instance)


def deregister(fn_plugin):
    global _override_instance
    if _override_instance is not None:
        omr.MRenderer.deregisterOverride(_override_instance)
        _override_instance = None
    fn_plugin.deregisterCommand(RenderSequenceCommand.kCmdName)