studio/rigging/sidecar.py

"""studio.rigging.sidecar — operator-annotation JSON loader + saver.

Sidecar format (versioned, forwards-compatible):

    {
        "schema_version": 1,
        "image": "goblin.png",
        "image_shape": [256, 256],
        "joints": {
            "head":        {"y": 50,  "x": 128, "confidence": 1.0},
            "neck":        {"y": 110, "x": 128},
            ...
        }
    }

Sidecar path convention: `<sprite>.joints.json` next to `<sprite>.png`.

Loader is lenient: missing joints default to image-center with confidence 0.0,
so the overlay renders them clearly out-of-place — the operator drags them to
the correct positions by editing the JSON.
"""
from __future__ import annotations

import json
from pathlib import Path
from typing import Optional, Tuple

import numpy as np

from pixel_cursor.rigging import JOINT_INDEX, JOINT_NAMES, NUM_JOINTS, Skeleton


SCHEMA_VERSION: int = 1


def sidecar_path_for(image_path: Path | str) -> Path:
    p = Path(image_path)
    return p.with_suffix(".joints.json")


def load_sidecar(
    image_path: Path | str,
    *,
    image_shape: Optional[Tuple[int, int]] = None,
) -> Optional[Skeleton]:
    """Return a Skeleton from the sidecar next to image_path, or None if absent.

    Missing or malformed joint entries default to image center with confidence 0.
    """
    sidecar = sidecar_path_for(image_path)
    if not sidecar.exists():
        return None

    data = json.loads(sidecar.read_text())
    version = int(data.get("schema_version", 1))
    if version > SCHEMA_VERSION:
        raise ValueError(
            f"sidecar schema_version {version} > supported {SCHEMA_VERSION} "
            f"({sidecar})"
        )

    shape_field = data.get("image_shape")
    if shape_field is not None:
        H, W = int(shape_field[0]), int(shape_field[1])
    elif image_shape is not None:
        H, W = image_shape
    else:
        raise ValueError(
            f"sidecar {sidecar} has no image_shape; pass image_shape= to load_sidecar"
        )

    joints_field = data.get("joints", {})
    positions = np.zeros((NUM_JOINTS, 2), dtype=np.float32)
    confidence = np.zeros(NUM_JOINTS, dtype=np.float32)

    for name in JOINT_NAMES:
        entry = joints_field.get(name)
        idx = JOINT_INDEX[name]
        if entry is None:
            positions[idx] = (H / 2.0, W / 2.0)
            confidence[idx] = 0.0
            continue
        positions[idx, 0] = float(entry["y"])
        positions[idx, 1] = float(entry["x"])
        confidence[idx] = float(entry.get("confidence", 1.0))

    return Skeleton(positions=positions, confidence=confidence, image_shape=(H, W))


def save_sidecar(
    image_path: Path | str,
    skeleton: Skeleton,
    *,
    image_name: Optional[str] = None,
) -> Path:
    sidecar = sidecar_path_for(image_path)
    name = image_name if image_name is not None else Path(image_path).name
    payload = {
        "schema_version": SCHEMA_VERSION,
        "image": name,
        "image_shape": list(skeleton.image_shape),
        "joints": {
            joint_name: {
                "y": round(float(skeleton.positions[i, 0]), 2),
                "x": round(float(skeleton.positions[i, 1]), 2),
                "confidence": round(float(skeleton.confidence[i]), 3),
            }
            for i, joint_name in enumerate(JOINT_NAMES)
        },
    }
    sidecar.write_text(json.dumps(payload, indent=2, sort_keys=False) + "\n")
    return sidecar


def sidecar_template(
    image_shape: Tuple[int, int],
    image_name: str = "",
) -> dict:
    """Return a template dict the operator can write to disk and fill in."""
    H, W = image_shape
    cy, cx = H / 2.0, W / 2.0
    return {
        "schema_version": SCHEMA_VERSION,
        "image": image_name,
        "image_shape": [H, W],
        "joints": {
            name: {"y": cy, "x": cx, "confidence": 0.0}
            for name in JOINT_NAMES
        },
    }


__all__ = [
    "SCHEMA_VERSION",
    "sidecar_path_for",
    "load_sidecar",
    "save_sidecar",
    "sidecar_template",
]