studio/io.py

"""I/O helpers: FrameStack <-> .npz round-trip, ndarray -> video file.

GIF export uses Pillow (always available). MP4/WebM/MKV export uses imageio
when installed (optional extra: `pip install 'venture-studio[video]'`); without
it, those extensions raise a clean ImportError with the install instruction.
"""
from __future__ import annotations
from pathlib import Path

import numpy as np
from PIL import Image as PILImage

from pixel_cursor import FrameStack
from pixel_cursor.artifact import _new_framestack


def save_framestack_npz(stack: FrameStack, path: Path | str) -> Path:
    out = Path(path)
    np.savez_compressed(
        str(out),
        frames=stack.frames,
        fps=np.array(stack.fps, dtype=np.int32),
        name=np.array(stack.name, dtype="U256"),
    )
    return out


def load_framestack_npz(path: Path | str) -> FrameStack:
    data = np.load(str(path), allow_pickle=False)
    name = str(data["name"]) if "name" in data.files else ""
    return _new_framestack(data["frames"], fps=int(data["fps"]), name=name)


def export_video(stack: FrameStack, out_path: Path | str) -> Path:
    out = Path(out_path)
    ext = out.suffix.lower()
    if ext == ".gif":
        return _export_gif(stack, out)
    if ext in (".mp4", ".webm", ".mkv", ".avi"):
        return _export_imageio(stack, out)
    raise ValueError(
        f"Unsupported export extension {ext!r}; supported: .gif .mp4 .webm .mkv .avi"
    )


def _export_gif(stack: FrameStack, out: Path) -> Path:
    frames = [PILImage.fromarray(stack.frames[i]) for i in range(stack.num_frames)]
    duration_ms = max(1, 1000 // max(1, stack.fps))
    frames[0].save(
        out, save_all=True, append_images=frames[1:], duration=duration_ms, loop=0
    )
    return out


def _export_imageio(stack: FrameStack, out: Path) -> Path:
    try:
        import imageio.v3 as iio
    except ImportError as e:
        raise ImportError(
            f"Exporting to {out.suffix} requires imageio. Install with: "
            f"pip install 'venture-studio[video]'  (or: pip install imageio imageio-ffmpeg)"
        ) from e
    iio.imwrite(str(out), stack.frames, fps=stack.fps)
    return out