fdanyone/skeleton/render_worker.py

"""CPU-only target skeleton renderer for conditioning overlap."""

from __future__ import annotations

import json
import os
import sys
import time
from fractions import Fraction
from pathlib import Path

import numpy as np

from fdanyone.errors import FourDAnyoneError
from fdanyone.geometry.cameras import Camera
from fdanyone.io import write_json
from fdanyone.skeleton.pipeline import render_skeleton_video


def _camera_from_payload(payload: object) -> Camera:
    """Decode a camera record from the conditioning file protocol."""

    if not isinstance(payload, dict):
        raise FourDAnyoneError("Target render camera must be a JSON object.")
    return Camera.from_dict(payload)


def render_target_skeletons(request_path: str | Path) -> tuple[Path, ...]:
    """Render every target in one prepared, CUDA-hidden file request."""

    started: float = time.monotonic()
    request_file: Path = Path(request_path).expanduser().resolve()
    request: object = json.loads(request_file.read_text())
    if not isinstance(request, dict) or int(request.get("schema_version", 0)) != 1:
        raise FourDAnyoneError("Unsupported target skeleton render request.")
    cuda_visible_devices: str | None = os.environ.get("CUDA_VISIBLE_DEVICES")
    if cuda_visible_devices not in ("", "-1"):
        raise FourDAnyoneError(
            "Target skeleton rendering must run with CUDA_VISIBLE_DEVICES disabled."
        )

    keypoints_path: Path = Path(str(request["keypoints_3d_path"])).expanduser().resolve()
    keypoints_world_fkc: np.ndarray = np.load(keypoints_path, allow_pickle=False)
    if keypoints_world_fkc.dtype != np.float32 or keypoints_world_fkc.ndim != 3:
        raise FourDAnyoneError(
            "Prepared target keypoints must be float32 [frames,keypoints,3], got "
            f"dtype={keypoints_world_fkc.dtype}, shape={keypoints_world_fkc.shape}."
        )
    targets: object = request.get("targets")
    if not isinstance(targets, list) or not targets:
        raise FourDAnyoneError("Target skeleton render request has no targets.")

    rendered_paths: list[Path] = []
    target: object
    for target in targets:
        if not isinstance(target, dict):
            raise FourDAnyoneError("Target skeleton entry must be a JSON object.")
        camera: Camera = _camera_from_payload(target["camera"])
        output_path: Path = Path(str(target["output_path"])).expanduser().resolve()
        rendered_path: Path = render_skeleton_video(
            keypoints_world_fkc=keypoints_world_fkc,
            camera=camera,
            output_path=output_path,
            num_frames=int(request["num_frames"]),
            canvas_height=int(request["canvas_height"]),
            canvas_width=int(request["canvas_width"]),
            output_height=int(request["output_height"]),
            output_width=int(request["output_width"]),
            body_height=float(request["body_height"]),
            focal_pixels=float(request["focal_pixels"]),
            fps=Fraction(int(request["fps_num"]), int(request["fps_den"])),
            crf=int(request["crf"]),
            preset=str(request["preset"]),
        )
        rendered_paths.append(rendered_path)

    done_path: Path = Path(str(request["done_path"])).expanduser().resolve()
    write_json(
        done_path,
        {
            "schema_version": 1,
            "rendered_views": len(rendered_paths),
            "elapsed_seconds": time.monotonic() - started,
            "cuda_visible_devices": cuda_visible_devices,
            "outputs": [str(path) for path in rendered_paths],
        },
    )
    return tuple(rendered_paths)


def main(request_path: str) -> None:
    """Run one target skeleton request."""

    render_target_skeletons(request_path)


if __name__ == "__main__":
    if len(sys.argv) != 2:
        raise SystemExit("Usage: python -m fdanyone.skeleton.render_worker REQUEST.json")
    main(sys.argv[1])