studio/animation.py

"""Headless learned-rig retargeting, rendered evidence and full-frame collision review."""

import json
import os
from pathlib import Path
import shutil
import subprocess

ROOT = Path(__file__).resolve().parents[1]


def _run(command, log_path, timeout=1800):
    with Path(log_path).open("w") as log:
        result = subprocess.run(command, stdout=log, stderr=subprocess.STDOUT, timeout=timeout)
    if result.returncode:
        tail = Path(log_path).read_text(errors="replace")[-4000:]
        raise RuntimeError(f"Animation stage failed ({result.returncode}): {tail}")


def animate_character(
    rigged_glb, motion_npz, motion_bvh, output_dir, *, body_clearance=True, smooth_joints=False
):
    """Retarget real Kimodo motion to a learned Mixamo rig; never substitutes a canned clip."""
    import numpy as np

    out = Path(output_dir).resolve()
    out.mkdir(parents=True, exist_ok=True)
    blender = os.getenv("BLENDER_BIN") or shutil.which("blender")
    if not blender and Path("/Applications/Blender.app/Contents/MacOS/Blender").exists():
        blender = "/Applications/Blender.app/Contents/MacOS/Blender"
    if not blender:
        raise RuntimeError("Blender is required for animation export and collision auditing.")
    ffmpeg = shutil.which("ffmpeg")
    if not ffmpeg:
        raise RuntimeError("ffmpeg is required for the downloadable animation video.")
    with np.load(motion_npz, allow_pickle=False) as motion:
        poses = motion["posed_joints"]
        if poses.ndim != 3 or poses.shape[2] != 3 or len(poses) < 2 or not np.isfinite(poses).all():
            raise ValueError("Kimodo output must contain finite frame/joint/xyz motion.")
        frames = len(poses)

    def blender_stage(script, arguments, name, timeout=1800):
        _run(
            [
                blender,
                "--background",
                "--python-exit-code",
                "1",
                "--python",
                str(ROOT / "scripts" / script),
                "--",
                *map(str, arguments),
            ],
            out / f"{name}.log",
            timeout,
        )

    def retarget(destination, smooth, name):
        destination.mkdir(parents=True, exist_ok=True)
        blender_stage(
            "retarget_character.py",
            [
                "--input",
                Path(rigged_glb).resolve(),
                "--motion",
                Path(motion_npz).resolve(),
                "--bvh",
                Path(motion_bvh).resolve(),
                "--output",
                destination,
                *(["--body-clearance"] if body_clearance else []),
                *(["--smooth-joints", "--arm-blend-band", "0.9"] if smooth else []),
            ],
            name,
        )
        blender_stage(
            "audit_character_collisions.py",
            ["--input", destination / "animation.glb", "--output", destination / "audit", "--frames", frames],
            name + "-audit",
            3600,
        )
        audit = json.loads((destination / "audit" / "audit.json").read_text())
        technical_animation_check(audit, frames)
        (destination / "audit" / "audit.json").write_text(json.dumps(audit, indent=2))
        return audit

    selected = out
    audit = retarget(out, smooth_joints, "retarget")
    selection = {"selected": "baseline", "baseline": audit["summary"]}
    if not audit["summary"]["technical_pass"] and not smooth_joints:
        candidate_dir = out / "repairs" / "smooth-joints"
        candidate = retarget(candidate_dir, True, "smooth-joints")
        improved, comparison = repair_improves(audit, candidate, frames)
        selection.update(candidate=candidate["summary"], comparison=comparison)
        if improved:
            selected, audit = candidate_dir, candidate
            selection["selected"] = "smooth-joints"
    (out / "repair-selection.json").write_text(json.dumps(selection, indent=2))
    glb = selected / "animation.glb"
    blender_stage(
        "render_animation.py",
        [
            "--input",
            glb,
            "--output",
            out / "frames",
            "--step",
            1,
            "--size",
            512,
            "--angle",
            "front",
            "--frame-count",
            frames,
        ],
        "render",
        3600,
    )
    video = out / "animation.mp4"
    _run(
        [
            ffmpeg,
            "-y",
            "-framerate",
            "30",
            "-i",
            str(out / "frames" / "%04d.png"),
            "-c:v",
            "libx264",
            "-pix_fmt",
            "yuv420p",
            "-movflags",
            "+faststart",
            str(video),
        ],
        out / "video.log",
    )
    from PIL import Image, ImageOps, ImageDraw

    samples = sorted(set(round(i * (frames - 1) / 15) + 1 for i in range(16)))
    sheet = Image.new("RGB", (4 * 256, 4 * 280), "#182830")
    draw = ImageDraw.Draw(sheet)
    for index, frame in enumerate(samples):
        with Image.open(out / "frames" / f"{frame:04d}.png") as image:
            sheet.paste(ImageOps.contain(image, (256, 256)), ((index % 4) * 256, (index // 4) * 280))
        draw.text(((index % 4) * 256 + 8, (index // 4) * 280 + 258), f"Frame {frame}", fill="white")
    sheet.save(out / "contact-sheet.jpg")
    blender_stage(
        "render_animation.py",
        [
            "--input",
            glb,
            "--output",
            out / "side",
            "--preview",
            "--angle",
            "side",
            "--size",
            320,
            "--frame-count",
            frames,
        ],
        "side-render",
    )
    side_files = sorted((out / "side").glob("*.png"))
    side_sheet = Image.new("RGB", (320 * len(side_files), 320), "#182830")
    for index, path in enumerate(side_files):
        with Image.open(path) as image:
            side_sheet.paste(image, (index * 320, 0))
    side_sheet.save(out / "side-contact-sheet.jpg")
    audit = json.loads((selected / "audit" / "audit.json").read_text())
    technical_pass = technical_animation_check(audit, frames)
    (selected / "audit" / "audit.json").write_text(json.dumps(audit, indent=2))
    return {
        "animated_glb": str(glb),
        "video": str(video),
        "blend": str(selected / "animation.blend"),
        "retargeting": str(selected / "retargeting.json"),
        "repair_selection": str(out / "repair-selection.json"),
        "contact_sheet": str(out / "contact-sheet.jpg"),
        "contact_sheets": [str(out / "contact-sheet.jpg"), str(out / "side-contact-sheet.jpg")],
        "technical_pass": technical_pass,
        "audit": str(selected / "audit" / "audit.json"),
        "frames": frames,
        "fps": 30,
        "collision_summary": audit["summary"],
    }


def technical_animation_check(audit, frames):
    import numpy as np

    tracks = np.array([[row[key] for key in sorted(row)] for row in audit["joints"]])
    movement = float(np.linalg.norm(tracks - tracks[:1], axis=2).max())
    summary = audit["summary"]
    technical_pass = bool(
        summary["frames"] == frames
        and len(tracks) == frames
        and summary["peak_new_pairs"] == 0
        and summary["ground_min_z"] >= -1e-5
        and summary.get("ground_min_all_vertices_z", summary["ground_min_z"]) >= -1e-5
        and np.isfinite(tracks).all()
        and movement > 0.01
    )
    summary["maximum_joint_displacement"] = movement
    summary["technical_pass"] = technical_pass
    return technical_pass


def repair_improves(baseline, candidate, frames):
    """Choose a repair only when geometry improves without sacrificing sampled motion."""
    import numpy as np

    def tracks(audit):
        names = sorted(audit["joints"][0])
        values = np.array([[row[name] for name in names] for row in audit["joints"]])
        anchor = names.index("Head") if "Head" in names else 0
        return values - values[:, anchor : anchor + 1], float(
            np.linalg.norm(values - values[:1], axis=2).max()
        )

    (before, before_motion), (after, after_motion) = tracks(baseline), tracks(candidate)
    compatible = before.shape == after.shape and len(after) == frames
    finite = bool(np.isfinite(after).all())
    difference = float(np.linalg.norm(after - before, axis=2).max()) if compatible else float("inf")
    left, right = baseline["summary"], candidate["summary"]
    checks = {
        "all_frames_finite": compatible and finite,
        "motion_preserved": compatible
        and finite
        and difference <= 1e-4
        and after_motion >= max(0.01, before_motion * 0.95),
        "maximum_relative_joint_difference": difference,
        "fewer_intersections": right["peak_new_pairs"] < left["peak_new_pairs"],
        "no_rest_regression": right.get("rest_intersections", 0) <= left.get("rest_intersections", 0),
        "no_deeper_crossings": right["peak_triangle_plane_depth"] <= left["peak_triangle_plane_depth"] + 1e-6,
        "no_more_affected_surface": right["peak_affected_surface_fraction"]
        <= left["peak_affected_surface_fraction"] + 1e-6,
        "no_ground_regression": right["ground_min_all_vertices_z"]
        >= min(left["ground_min_all_vertices_z"], -1e-5) - 1e-6,
    }
    accepted = all(value for key, value in checks.items() if key != "maximum_relative_joint_difference")
    return bool(accepted), checks