fn_registry/python/functions/ml/render_openpose_walk_skeletons.py

"""Dibuja una secuencia de esqueletos OpenPose (COCO-18) de un ciclo de caminar.

Funcion impura: rasteriza con PIL (Pillow) N frames de un walk cycle lateral y
escribe cada uno como PNG sobre fondo negro. Cada PNG es una pose-map valida para
el ControlNet OpenPose de SD1.5 (control_v11p_sd15_openpose_fp16): el esqueleto NO
lo genera la IA, lo aportas tu dibujado y el modelo anima el personaje sobre el.

El formato es el render clasico de OpenPose: 18 keypoints (COCO-18), 17 limbs,
colores canonicos por articulacion/par. Las piernas oscilan en oposicion a los
brazos y el cuerpo "bota" (sube en passing, baja en contact) — un walk cycle de
manual de animacion. Para frames=4 produce las 4 fases canonicas
(contact-izq, passing, contact-der, passing); para mas frames muestrea el ciclo
parametrico de forma continua (interpolando las fases intermedias down/up).
"""
import json
import math
import os


# Orden de indice COCO-18 que espera el ControlNet OpenPose.
COCO18_NAMES = [
    "nose", "neck", "right_shoulder", "right_elbow", "right_wrist",
    "left_shoulder", "left_elbow", "left_wrist", "right_hip", "right_knee",
    "right_ankle", "left_hip", "left_knee", "left_ankle",
    "right_eye", "left_eye", "right_ear", "left_ear",
]

# Pares (limbs) — formato OpenPose clasico (17 limbs sobre 18 keypoints).
LIMB_SEQ = [
    (1, 2), (1, 5), (2, 3), (3, 4), (5, 6), (6, 7), (1, 8), (8, 9), (9, 10),
    (1, 11), (11, 12), (12, 13), (1, 0), (0, 14), (14, 16), (0, 15), (15, 17),
]

# Colores canonicos (RGB) por indice, identicos a los del render original de
# OpenPose / controlnet_aux. El keypoint i usa COLORS[i]; el limb i usa COLORS[i].
COLORS = [
    [255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0],
    [85, 255, 0], [0, 255, 0], [0, 255, 85], [0, 255, 170], [0, 255, 255],
    [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255], [170, 0, 255],
    [255, 0, 255], [255, 0, 170], [255, 0, 85],
]


def _walk_pose_norm(t: float) -> dict:
    """Devuelve los 18 keypoints normalizados (0..1, y hacia abajo) de un walk
    cycle lateral mirando a la derecha, para la fase t in [0, 1).

    t=0.0 -> contact (pierna izq adelante), t=0.25 -> passing (cuerpo arriba),
    t=0.5 -> contact (pierna der adelante), t=0.75 -> passing.
    """
    two_pi = 2.0 * math.pi
    cx = 0.50
    neck_y = 0.245
    sh_y = 0.265
    hip_y0 = 0.55
    ground_y = 0.885
    thigh_drop = 0.165   # caida vertical hip->knee
    stride = 0.105       # desplazamiento horizontal max del pie respecto al hip
    bob_amp = 0.025      # rebote vertical del cuerpo
    lift_amp = 0.05      # cuanto se levanta el pie en swing
    sh_dx = 0.024        # media separacion horizontal de hombros (profundidad)
    hip_dx = 0.026       # media separacion horizontal de caderas

    # Rebote: bajo en contact (t=0, .5), alto en passing (t=.25, .75).
    rise = bob_amp * (1.0 - math.cos(2.0 * two_pi * t)) / 2.0
    nx = cx
    ny = neck_y - rise
    hip_cy = hip_y0 - rise

    pts = {}
    pts[1] = (nx, ny)                        # neck
    pts[0] = (nx + 0.055, ny - 0.105)        # nose (mira a la derecha)
    pts[14] = (nx + 0.045, ny - 0.120)       # right_eye
    pts[15] = (nx + 0.026, ny - 0.120)       # left_eye
    pts[16] = (nx + 0.002, ny - 0.103)       # right_ear (atras)
    pts[17] = (nx - 0.012, ny - 0.097)       # left_ear (atras)
    pts[2] = (nx - sh_dx, sh_y - rise)       # right_shoulder
    pts[5] = (nx + sh_dx, sh_y - rise)       # left_shoulder

    r_hip = (cx - hip_dx, hip_cy)
    l_hip = (cx + hip_dx, hip_cy)
    pts[8] = r_hip
    pts[11] = l_hip

    # Factores de oscilacion de las piernas (opuestas entre si).
    fwd_l = math.cos(two_pi * t)             # pierna izq adelante en t=0
    fwd_r = -fwd_l                           # pierna der opuesta
    lift_l = lift_amp * max(0.0, -math.sin(two_pi * t))  # izq levanta t in (.5,1)
    lift_r = lift_amp * max(0.0, math.sin(two_pi * t))   # der levanta t in (0,.5)

    def _leg(hip, fwd, lift):
        hx, hy = hip
        ax = hx + stride * fwd
        ay = ground_y - lift
        bend = 0.012 + 0.06 * (lift / lift_amp if lift_amp else 0.0)
        kx = (hx + ax) / 2.0 + bend          # rodilla apunta hacia delante
        ky = hy + thigh_drop
        return (kx, ky), (ax, ay)

    rk, ra = _leg(r_hip, fwd_r, lift_r)
    lk, la = _leg(l_hip, fwd_l, lift_l)
    pts[9], pts[10] = rk, ra                 # right_knee, right_ankle
    pts[12], pts[13] = lk, la                # left_knee, left_ankle

    # Brazos en oposicion: brazo der adelante cuando pierna izq adelante.
    arm_fwd_r = fwd_l
    arm_fwd_l = fwd_r
    sh_to_el = 0.105
    el_to_wr = 0.110
    arm_sw = 0.05

    def _arm(sh, fwd):
        sx, sy = sh
        ex = sx + arm_sw * fwd + 0.008
        ey = sy + sh_to_el
        wx = sx + arm_sw * 1.7 * fwd + 0.016
        wy = ey + el_to_wr
        return (ex, ey), (wx, wy)

    re, rw = _arm(pts[2], arm_fwd_r)
    le, lw = _arm(pts[5], arm_fwd_l)
    pts[3], pts[4] = re, rw                  # right_elbow, right_wrist
    pts[6], pts[7] = le, lw                  # left_elbow, left_wrist

    return pts


def render_openpose_walk_skeletons(
    out_dir: str,
    *,
    frames: int = 4,
    width: int = 512,
    height: int = 768,
    facing: str = "right",
    line_width: int = 4,
    point_radius: int = 6,
    filename_prefix: str = "walk_pose",
) -> dict:
    """Dibuja N esqueletos OpenPose COCO-18 de un walk cycle y los guarda como PNG.

    Args:
        out_dir: directorio destino de los PNG; se crea si no existe.
        frames: numero de fases del ciclo a renderizar (default 4 = las 4
            fases canonicas contact/passing/contact/passing). keyword-only.
        width: ancho de cada PNG en pixeles (default 512). keyword-only.
        height: alto de cada PNG en pixeles (default 768, retrato). keyword-only.
        facing: "right" (mira a +x) o "left" (espeja en X). keyword-only.
        line_width: grosor en px de las lineas de los limbs (default 4).
            keyword-only.
        point_radius: radio en px de los circulos de cada keypoint (default 6).
            keyword-only.
        filename_prefix: prefijo de los archivos; se nombran
            "<prefix>_<NN>.png" con NN de dos digitos. keyword-only.

    Returns:
        dict con:
        - ok (bool): True si todos los PNG se generaron.
        - skeleton_paths (list[str]): rutas de los PNG creados, en orden de fase.
        - frames (int): numero de frames generados.
        - width (int): ancho usado.
        - height (int): alto usado.
        - error (str): mensaje de error; cadena vacia si todo OK.

    No lanza salvo que out_dir sea None: cualquier otro fallo se captura en
    el campo "error" con ok=False.
    """
    if out_dir is None:
        raise ValueError("out_dir no puede ser None")

    out = {
        "ok": False, "skeleton_paths": [], "frames": int(frames or 0),
        "width": int(width or 0), "height": int(height or 0), "error": "",
    }

    try:
        from PIL import Image, ImageDraw
    except ImportError:
        out["error"] = "PIL (Pillow) no esta instalado en este interprete"
        return out

    try:
        frames = int(frames)
        width = int(width)
        height = int(height)
        line_width = max(1, int(line_width))
        point_radius = max(1, int(point_radius))
    except (TypeError, ValueError) as exc:
        out["error"] = f"argumento numerico invalido: {exc}"
        return out

    if frames <= 0:
        out["error"] = "frames debe ser >= 1"
        return out
    if width < 16 or height < 16:
        out["error"] = "width y height deben ser >= 16"
        return out
    if facing not in ("right", "left"):
        out["error"] = f"facing debe ser 'right' o 'left', no {facing!r}"
        return out

    out["frames"] = frames
    out["width"] = width
    out["height"] = height

    try:
        os.makedirs(out_dir, exist_ok=True)
    except OSError as exc:
        out["error"] = f"no se pudo crear out_dir {out_dir!r}: {exc}"
        return out

    paths = []
    try:
        for i in range(frames):
            t = i / float(frames)
            norm = _walk_pose_norm(t)

            # Mirror en X para facing izquierda (el esqueleto sigue siendo valido).
            kp = {}
            for idx, (x, y) in norm.items():
                if facing == "left":
                    x = 1.0 - x
                kp[idx] = (x * width, y * height)

            img = Image.new("RGB", (width, height), (0, 0, 0))
            draw = ImageDraw.Draw(img)

            # Limbs primero (lineas gruesas del color del par).
            for li, (a, b) in enumerate(LIMB_SEQ):
                if a in kp and b in kp:
                    col = tuple(COLORS[li % len(COLORS)])
                    draw.line([kp[a], kp[b]], fill=col, width=line_width)

            # Keypoints encima (circulos rellenos del color de la articulacion).
            for idx in range(len(COCO18_NAMES)):
                if idx not in kp:
                    continue
                cx, cy = kp[idx]
                col = tuple(COLORS[idx % len(COLORS)])
                draw.ellipse(
                    [cx - point_radius, cy - point_radius,
                     cx + point_radius, cy + point_radius],
                    fill=col,
                )

            path = os.path.join(out_dir, f"{filename_prefix}_{i:02d}.png")
            img.save(path)
            paths.append(path)
    except (OSError, ValueError) as exc:
        out["error"] = f"fallo al rasterizar/guardar el frame {len(paths)}: {exc}"
        out["skeleton_paths"] = paths
        return out

    out["ok"] = True
    out["skeleton_paths"] = paths
    return out


if __name__ == "__main__":
    res = render_openpose_walk_skeletons("/tmp/walk_skeletons_demo", frames=4)
    print(json.dumps(res, indent=2))