ovi/utils/av_edit_data.py

"""Shared data and media utilities for audio-video editing."""

from __future__ import annotations

import json
import math
from pathlib import Path
from typing import Any, Mapping, Optional

import imageio
import librosa
import numpy as np
import pandas as pd
import torch
import torchvision
from PIL import Image
from scipy.io import wavfile


CANONICAL_COLUMNS = (
    "source_video",
    "source_audio",
    "target_video",
    "target_audio",
    "instruction",
)


def load_manifest(path: str | Path) -> list[dict[str, Any]]:
    """Load a CSV, JSON, or JSONL manifest into row dictionaries."""

    manifest_path = Path(path).expanduser().resolve()
    if not manifest_path.is_file():
        raise FileNotFoundError(f"Manifest does not exist: {manifest_path}")

    suffix = manifest_path.suffix.lower()
    if suffix == ".csv":
        frame = pd.read_csv(manifest_path)
        return [frame.iloc[index].to_dict() for index in range(len(frame))]
    if suffix == ".jsonl":
        rows = []
        with manifest_path.open("r", encoding="utf-8") as handle:
            for line_number, line in enumerate(handle, start=1):
                line = line.strip()
                if not line:
                    continue
                value = json.loads(line)
                if not isinstance(value, dict):
                    raise ValueError(
                        f"JSONL row {line_number} must be an object: {manifest_path}"
                    )
                rows.append(value)
        return rows
    if suffix == ".json":
        with manifest_path.open("r", encoding="utf-8") as handle:
            value = json.load(handle)
        if not isinstance(value, list) or not all(isinstance(row, dict) for row in value):
            raise ValueError(f"JSON manifest must contain a list of objects: {manifest_path}")
        return value

    raise ValueError(f"Unsupported manifest format '{suffix}'. Use CSV, JSON, or JSONL.")


def is_missing(value: Any) -> bool:
    if value is None:
        return True
    if isinstance(value, str):
        return not value.strip()
    try:
        return bool(pd.isna(value))
    except (TypeError, ValueError):
        return False


def get_instruction(
    row: Mapping[str, Any], instruction_column: str = "instruction"
) -> str:
    """Read one non-empty editing instruction from a manifest row."""

    if instruction_column not in row:
        raise KeyError(
            f"Manifest is missing instruction column '{instruction_column}'."
        )
    value = row[instruction_column]
    if is_missing(value):
        raise ValueError(
            f"Instruction column '{instruction_column}' contains an empty value."
        )
    return str(value).strip()


def resolve_media_path(
    value: Any,
    manifest_path: str | Path,
    base_path: Optional[str | Path] = None,
    required: bool = True,
) -> Optional[Path]:
    """Resolve a manifest media path relative to base_path or the manifest directory."""

    if is_missing(value):
        if required:
            raise ValueError("A required media path is empty.")
        return None

    path = Path(str(value)).expanduser()
    if not path.is_absolute():
        root = Path(base_path).expanduser() if base_path else Path(manifest_path).parent
        path = root / path
    path = path.resolve()
    if not path.is_file():
        raise FileNotFoundError(f"Media file does not exist: {path}")
    return path


def get_video_info(path: str | Path) -> tuple[float, int]:
    reader = imageio.get_reader(str(path))
    try:
        metadata = reader.get_meta_data()
        fps = float(metadata.get("fps", 24.0))
        frame_count = int(reader.count_frames())
    finally:
        reader.close()
    if fps <= 0 or frame_count <= 0:
        raise ValueError(f"Invalid video metadata for {path}: fps={fps}, frames={frame_count}")
    return fps, frame_count


def snap_num_frames(frame_count: int, factor: int = 4, remainder: int = 1) -> int:
    while frame_count > 1 and frame_count % factor != remainder:
        frame_count -= 1
    if frame_count < 5:
        raise ValueError(f"At least 5 usable frames are required, got {frame_count}.")
    return frame_count


def _target_size(
    image: Image.Image,
    height: Optional[int],
    width: Optional[int],
    max_pixels: int,
    division_factor: int = 16,
) -> tuple[int, int]:
    if (height is None) != (width is None):
        raise ValueError("height and width must either both be set or both be null.")
    if height is not None and width is not None:
        return int(height), int(width)

    image_width, image_height = image.size
    if image_width * image_height > max_pixels:
        scale = math.sqrt((image_width * image_height) / max_pixels)
        image_height = int(image_height / scale)
        image_width = int(image_width / scale)
    image_height = max(division_factor, image_height // division_factor * division_factor)
    image_width = max(division_factor, image_width // division_factor * division_factor)
    return image_height, image_width


def crop_and_resize(image: Image.Image, height: int, width: int) -> Image.Image:
    image_width, image_height = image.size
    scale = max(width / image_width, height / image_height)
    image = torchvision.transforms.functional.resize(
        image,
        (round(image_height * scale), round(image_width * scale)),
        interpolation=torchvision.transforms.InterpolationMode.BILINEAR,
    )
    return torchvision.transforms.functional.center_crop(image, (height, width))


def load_video_array(
    path: str | Path,
    num_frames: int,
    height: Optional[int],
    width: Optional[int],
    max_pixels: int,
    target_size: Optional[tuple[int, int]] = None,
) -> tuple[np.ndarray, tuple[int, int]]:
    reader = imageio.get_reader(str(path))
    try:
        frames = []
        resolved_size = target_size
        for frame_index in range(num_frames):
            frame = Image.fromarray(reader.get_data(frame_index)).convert("RGB")
            if resolved_size is None:
                resolved_size = _target_size(frame, height, width, max_pixels)
            frame = crop_and_resize(frame, *resolved_size)
            frames.append(np.asarray(frame))
    finally:
        reader.close()

    video = np.stack(frames, axis=0).transpose(3, 0, 1, 2)
    return video, resolved_size


def load_audio_array(
    path: str | Path,
    sample_rate: int,
    num_samples: int,
) -> np.ndarray:
    audio = load_audio(path, sample_rate=sample_rate)
    return pad_or_trim_audio(audio, num_samples=num_samples)


def load_audio(path: str | Path, sample_rate: int) -> np.ndarray:
    """Load a complete mono audio stream without changing its duration."""

    audio, _ = librosa.load(str(path), sr=sample_rate, mono=True)
    return np.asarray(audio, dtype=np.float32)


def pad_or_trim_audio(audio: np.ndarray, num_samples: int) -> np.ndarray:
    """Pad or trim a waveform to an exact number of samples."""

    audio = np.asarray(audio, dtype=np.float32)
    if len(audio) >= num_samples:
        return audio[:num_samples]
    return np.pad(audio, (0, num_samples - len(audio)))


def save_audio(path: str | Path, audio: np.ndarray, sample_rate: int = 16000) -> None:
    output_path = Path(path)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    audio = np.clip(np.asarray(audio).squeeze(), -1.0, 1.0)
    wavfile.write(str(output_path), sample_rate, (audio * 32767).astype(np.int16))


def to_video_tensor(video: np.ndarray, device: Any, dtype: torch.dtype) -> torch.Tensor:
    tensor = torch.from_numpy(video).float().unsqueeze(0).to(device=device, dtype=dtype)
    return tensor / 127.5 - 1.0


def to_audio_tensor(audio: np.ndarray, device: Any) -> torch.Tensor:
    return torch.from_numpy(audio).float().unsqueeze(0).to(device)