fdanyone/runs.py

"""Read what one finished 4DAnyone inference run left on disk.

Two readers live here. ``discover_run`` collects every artifact of a clip --
motion, generated videos, skeletons, and the source clip -- into a
``RunLayout``. The rig readers parse the ``cameras.json`` that
``fdanyone.output`` writes next to the generated videos.
"""

from __future__ import annotations

import json
import logging
from dataclasses import dataclass
from pathlib import Path

from fdanyone.errors import FourDAnyoneError
from fdanyone.io import read_json

LOGGER = logging.getLogger("fdanyone.runs")


# ---------------------------------------------------------------------------
# Run layout
# ---------------------------------------------------------------------------


@dataclass(frozen=True)
class RunLayout:
    """Files that one inference run left on disk for a single clip."""

    clip: str
    data_dir: Path
    result_dir: Path
    motion_dir: Path
    cameras_json: Path | None
    metadata_json: Path | None
    source_video: Path | None
    dense_videos: tuple[tuple[int, Path], ...]
    skeleton_videos: tuple[tuple[int, Path], ...]
    sparse_videos: tuple[tuple[int, Path], ...]

    @property
    def has_motion(self) -> bool:
        return (self.motion_dir / "motion.json").is_file() and (self.motion_dir / "motion.safetensors").is_file()


def numbered_videos(directory: Path) -> tuple[tuple[int, Path], ...]:
    """Return ``(camera_id, path)`` pairs for ``NN.mp4`` files, ID-ordered."""

    if not directory.is_dir():
        return ()
    found: list[tuple[int, Path]] = []
    for path in sorted(directory.glob("*.mp4")):
        try:
            found.append((int(path.stem), path))
        except ValueError:
            LOGGER.warning("Ignoring video with a non-numeric name: %s", path)
    return tuple(sorted(found))


def find_source_video(data_dir: Path, clip: str, filename: str | None) -> Path | None:
    """Locate the input clip that produced the run."""

    candidates: list[Path] = []
    if filename:
        candidates.append(data_dir / "source" / "pexels" / filename)
    for suffix in (".mp4", ".mov", ".mkv", ".webm", ".MP4", ".MOV"):
        candidates.append(data_dir / "source" / "pexels" / f"{clip}{suffix}")
    for candidate in candidates:
        if candidate.is_file():
            return candidate
    # Only the input and result trees can hold a clip; the rest of the data
    # root is model output that a recursive walk would scan for nothing.
    roots = tuple(root for root in (data_dir / "source", data_dir / "fdanyone") if root.is_dir())
    if filename:
        matches = sorted(path for root in roots for path in root.rglob(filename) if path.is_file())
        if matches:
            return matches[0]
    matches = sorted(
        path for root in roots for path in root.rglob(f"{clip}.*") if path.is_file() and path.suffix != ".rrd"
    )
    return matches[0] if matches else None


def discover_run(data_dir: Path, clip: str) -> RunLayout:
    """Collect every artifact of ``clip`` and reject an empty run."""

    result_dir = data_dir / "fdanyone" / clip
    motion_dir = data_dir / "gvhmr" / "results" / clip
    metadata_json = result_dir / "metadata.json"
    cameras_json = result_dir / "cameras.json"
    metadata = read_json(metadata_json)
    filename = None
    if metadata is not None:
        filename = str(metadata.get("input", {}).get("filename") or "") or None

    layout = RunLayout(
        clip=clip,
        data_dir=data_dir,
        result_dir=result_dir,
        motion_dir=motion_dir,
        cameras_json=cameras_json if cameras_json.is_file() else None,
        metadata_json=metadata_json if metadata_json.is_file() else None,
        source_video=find_source_video(data_dir, clip, filename),
        dense_videos=numbered_videos(result_dir / "videos" / "dense"),
        skeleton_videos=numbered_videos(result_dir / "skeletons"),
        sparse_videos=numbered_videos(result_dir / "videos" / "sparse"),
    )
    if not layout.has_motion and not layout.dense_videos and layout.source_video is None:
        raise FourDAnyoneError(
            f"Nothing to visualize for clip {clip!r}. Expected at least one of:\n"
            f"  {motion_dir / 'motion.safetensors'} (GVHMR motion)\n"
            f"  {result_dir / 'videos' / 'dense'} (generated views)\n"
            f"  a source clip named {clip}.* under {data_dir}\n"
            "Run `python inference.py --video_path <clip>` first."
        )
    return layout


# ---------------------------------------------------------------------------
# Camera rig
# ---------------------------------------------------------------------------


def read_cameras(result: Path) -> dict:
    """Read the ``cameras.json`` a result directory must carry."""

    path = result / "cameras.json"
    try:
        cameras = json.loads(path.read_text())
    except (FileNotFoundError, json.JSONDecodeError) as exc:
        raise FourDAnyoneError(f"Cannot read {path}: {exc}") from exc
    if not isinstance(cameras, dict):
        raise FourDAnyoneError("cameras.json must contain a JSON object.")
    return cameras


def camera_records(rig: dict) -> list[dict]:
    """Validate an OPENCV rig and return its ID-ordered camera records."""

    if not isinstance(rig, dict) or rig.get("camera_model") != "OPENCV":
        raise FourDAnyoneError("cameras.json has no supported OPENCV camera rig.")
    cameras = rig.get("cameras")
    if not isinstance(cameras, list) or not cameras:
        raise FourDAnyoneError("Camera rig must contain at least one camera.")
    if [camera.get("camera_id") for camera in cameras if isinstance(camera, dict)] != list(range(len(cameras))):
        raise FourDAnyoneError("Camera rig must be ordered by camera ID.")
    return cameras


def dense_video_paths(result: Path, cameras: list[dict]) -> tuple[Path, ...]:
    """Return the generated video of every camera, in rig order."""

    paths = []
    for camera in cameras:
        camera_id = int(camera["camera_id"])
        relative = f"videos/dense/{camera_id:02d}.mp4"
        if camera.get("video") != relative:
            raise FourDAnyoneError(f"Camera {camera_id:02d} points to the wrong target video.")
        paths.append(result / relative)
    return tuple(paths)