PoseRenderer/python/openpose_renderer/rasterizer.py

"""Draws OpenPose_full-style skeleton images (black background, colored
limbs/joints) and writes them out via maya.api.OpenMaya.MImage.

This is a self-contained software rasterizer (no OpenGL/render-target
dependency) operating on a flat RGBA bytearray, so it can be driven equally
from a batch/offline context or from inside a render override. Confidence-0
keypoints (unmapped OpenPose_full slots) are simply not drawn, and any limb
whose either endpoint is confidence-0 is skipped -- matching how OpenPose
itself omits undetected points/limbs.
"""

import maya.api.OpenMaya as om

from . import constants

BYTES_PER_PIXEL = 4  # RGBA


def _new_buffer(width, height):
    buf = bytearray(width * height * BYTES_PER_PIXEL)
    for i in range(0, len(buf), BYTES_PER_PIXEL):
        buf[i + 3] = 255  # opaque black background
    return buf


def _set_pixel(buffer, width, height, x, y, color):
    if x < 0 or x >= width or y < 0 or y >= height:
        return
    offset = (y * width + x) * BYTES_PER_PIXEL
    buffer[offset] = color[0]
    buffer[offset + 1] = color[1]
    buffer[offset + 2] = color[2]
    buffer[offset + 3] = 255


def _draw_circle(buffer, width, height, cx, cy, radius, color):
    if radius <= 0:
        return
    x_min = max(0, int(cx - radius))
    x_max = min(width - 1, int(cx + radius))
    y_min = max(0, int(cy - radius))
    y_max = min(height - 1, int(cy + radius))
    radius_sq = radius * radius
    for y in range(y_min, y_max + 1):
        dy = y + 0.5 - cy
        for x in range(x_min, x_max + 1):
            dx = x + 0.5 - cx
            if dx * dx + dy * dy <= radius_sq:
                _set_pixel(buffer, width, height, x, y, color)


def _draw_line(buffer, width, height, x0, y0, x1, y1, thickness, color):
    if thickness <= 0:
        return
    half = max(thickness * 0.5, 0.5)
    x_min = max(0, int(min(x0, x1) - half))
    x_max = min(width - 1, int(max(x0, x1) + half))
    y_min = max(0, int(min(y0, y1) - half))
    y_max = min(height - 1, int(max(y0, y1) + half))

    dx = x1 - x0
    dy = y1 - y0
    length_sq = dx * dx + dy * dy
    half_sq = half * half

    if length_sq < 1e-9:
        _draw_circle(buffer, width, height, x0, y0, half, color)
        return

    for y in range(y_min, y_max + 1):
        py = y + 0.5
        for x in range(x_min, x_max + 1):
            px = x + 0.5
            t = ((px - x0) * dx + (py - y0) * dy) / length_sq
            t = max(0.0, min(1.0, t))
            proj_x = x0 + t * dx
            proj_y = y0 + t * dy
            dist_sq = (px - proj_x) ** 2 + (py - proj_y) ** 2
            if dist_sq <= half_sq:
                _set_pixel(buffer, width, height, x, y, color)


def draw_character(buffer, width, height, keypoints, radius_scale=1.0, thickness_scale=1.0,
                    body_model=constants.BODY_MODEL_BODY_25):
    """Draw one character's OpenPose skeleton into an existing RGBA buffer.

    ``keypoints``: list of length constants.TOTAL_KEYPOINTS of
    (pixel_x, pixel_y, confidence), as returned by
    ``projector.project_character``. ``body_model`` selects BODY_25 (all 25
    body/foot points) or COCO (18 points, no MidHip/feet, different limb
    connectivity) -- hands and face are unaffected either way.
    """
    pairs, colors, thickness_list = constants.full_limb_data(body_model)
    for pair_i, (a, b) in enumerate(pairs):
        xa, ya, ca = keypoints[a]
        xb, yb, cb = keypoints[b]
        if ca <= 0.0 or cb <= 0.0:
            continue
        thickness = thickness_list[pair_i] * thickness_scale
        _draw_line(buffer, width, height, xa, ya, xb, yb, thickness, colors[pair_i])

    active_indices = constants.active_point_indices(body_model)
    for index in active_indices:
        x, y, confidence = keypoints[index]
        if confidence <= 0.0:
            continue
        radius = constants.KEYPOINT_POINT_RADIUS[index] * radius_scale
        _draw_circle(buffer, width, height, x, y, radius, constants.KEYPOINT_POINT_COLORS[index])


def render_characters_to_buffer(characters, width, height):
    """``characters``: iterable of (keypoints, radius_scale, thickness_scale, body_model).

    Returns a flat RGBA bytearray of size width*height*4.
    """
    buffer = _new_buffer(width, height)
    for keypoints, radius_scale, thickness_scale, body_model in characters:
        draw_character(buffer, width, height, keypoints, radius_scale, thickness_scale, body_model)
    return buffer


def _flip_rows(buffer, width, height):
    """MImage.setPixels() treats row 0 as the bottom of the image (Maya's
    classic bottom-up convention); everything else here (and every PNG
    reader downstream) treats row 0 as the top. Flip row order once, right
    before handing pixels to MImage, so the written file is a normal
    top-down image.
    """
    row_bytes = width * BYTES_PER_PIXEL
    flipped = bytearray(len(buffer))
    for y in range(height):
        src = y * row_bytes
        dst = (height - 1 - y) * row_bytes
        flipped[dst:dst + row_bytes] = buffer[src:src + row_bytes]
    return flipped


def render_characters_to_file(characters, width, height, file_path, file_format="png"):
    """Render all characters and write the result to ``file_path``."""
    buffer = render_characters_to_buffer(characters, width, height)
    image = om.MImage()
    image.setPixels(_flip_rows(buffer, width, height), width, height)
    image.writeToFile(file_path, outputFormat=file_format)