adam/oasis_dataset.py

from __future__ import annotations

import json
import re
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

from PIL import Image


BINARY_ACTIONS = {
    "w", "a", "s", "d", "jump", "arrow_up", "arrow_left", "arrow_down",
    "arrow_right", "enter", "shift", "ctrl", "alt", "tab", "escape",
    "q", "e", "r", "f", "z", "x", "c", "v", "key_1", "key_2", "key_3",
    "key_4", "mouse_left", "mouse_middle", "mouse_right",
}
CONTINUOUS_ACTIONS = {"mouse_dx", "mouse_dy", "zoom"}
DERIVED_ACTIONS = {
    "move_x", "move_y", "right_mouse", "camera_active",
    "camera_yaw_delta_degrees", "camera_pitch_delta_degrees",
    "mouse_raw_dx", "mouse_raw_dy",
}
SUPPORTED_ACTIONS = BINARY_ACTIONS | CONTINUOUS_ACTIONS | DERIVED_ACTIONS
REQUIRED_CANONICAL_ACTIONS = {"w", "a", "s", "d", "jump"}
IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp", ".bmp"}
RECORDED_ACTIONS = BINARY_ACTIONS | CONTINUOUS_ACTIONS
METADATA_FIELDS = {
    "session_id", "session_started_at", "frame_index", "filename",
    "timestamp_seconds", "camera_encoding",
}


@dataclass(slots=True)
class OasisDatasetReport:
    dataset_folders: list[str] = field(default_factory=list)
    frames: int = 0
    metadata_rows: int = 0
    valid_rows: int = 0
    valid_transitions: int = 0
    sessions: int = 0
    resolution: str = ""
    capture_fps: float | None = None
    recommended_frame_gap: int | None = None
    native_ai_fps: float | None = None
    action_counts: dict[str, int] = field(default_factory=dict)
    idle_rows: int = 0
    errors: list[str] = field(default_factory=list)
    warnings: list[str] = field(default_factory=list)

    @property
    def ok(self) -> bool:
        return not self.errors


def _read_dataset_info(directory: Path) -> dict[str, Any]:
    """Read optional recorder metadata without making it a dataset requirement."""
    try:
        value = json.loads((directory / "dataset_info.json").read_text(encoding="utf-8"))
        return value if isinstance(value, dict) else {}
    except (OSError, json.JSONDecodeError):
        return {}


def _recommended_frame_gap(capture_fps: float, _camera_encoding: str) -> int:
    """Choose a playable horizon that targets about 12 genuine AI frames per second.

    This intentionally differs from the external trainer's older movement-focused
    heuristic. A 12–15 FPS recording should train at gap 1, rather than being
    slowed to a 3–5 FPS playable world before GPU speed is even considered.
    """
    target_ai_fps = 12.0
    return max(1, min(12, round(capture_fps / target_ai_fps)))


def oasis_pace(
    value: str | list[str] | tuple[str, ...], *, frame_gap: int,
) -> dict[str, float | int | None]:
    """Return portable pacing information derived from connected recorder metadata.

    A model only makes one genuine frame for every prediction horizon. Display
    interpolation can look smoother, but cannot make controls more responsive.
    """
    rates: list[float] = []
    camera_encodings: list[str] = []
    for directory in dataset_directories(value):
        info = _read_dataset_info(directory)
        try:
            rate = float(info.get("capture_fps"))
        except (TypeError, ValueError):
            continue
        if rate > 0:
            rates.append(rate)
            camera_encodings.append(str(info.get("camera_encoding", "legacy_pixels")))
    if not rates or len({round(rate, 6) for rate in rates}) != 1:
        return {"capture_fps": None, "recommended_frame_gap": None, "native_ai_fps": None}
    capture_fps = rates[0]
    camera_encoding = (
        "relative_degrees_v1"
        if "relative_degrees_v1" in camera_encodings
        else camera_encodings[0]
    )
    gap = max(1, int(frame_gap))
    return {
        "capture_fps": capture_fps,
        "recommended_frame_gap": _recommended_frame_gap(capture_fps, camera_encoding),
        "native_ai_fps": capture_fps / gap,
    }


def dataset_directories(value: str | list[str] | tuple[str, ...]) -> list[Path]:
    entries = value if isinstance(value, (list, tuple)) else str(value or "").split(";")
    directories: list[Path] = []
    for entry in entries:
        text = str(entry).strip().strip('"')
        if not text:
            continue
        path = Path(text).expanduser()
        if path not in directories:
            directories.append(path)
    return directories


def _numeric_frame_index(path: Path) -> int | None:
    match = re.search(r"frame_(\d+)", path.stem, re.I) or re.search(r"(\d+)", path.stem)
    return int(match.group(1)) if match else None


def validate_oasis_dataset(
    value: str | list[str] | tuple[str, ...], *, frame_gap: int = 1,
    verify_images: bool = True,
) -> OasisDatasetReport:
    report = OasisDatasetReport()
    frame_gap = max(1, int(frame_gap))
    pace = oasis_pace(value, frame_gap=frame_gap)
    report.capture_fps = pace["capture_fps"]  # type: ignore[assignment]
    report.recommended_frame_gap = pace["recommended_frame_gap"]  # type: ignore[assignment]
    report.native_ai_fps = pace["native_ai_fps"]  # type: ignore[assignment]
    directories = dataset_directories(value)
    if not directories:
        report.errors.append("Select at least one Oasis action dataset folder.")
        return report
    seen_resolution: tuple[int, int] | None = None
    action_counts = {name: 0 for name in sorted(RECORDED_ACTIONS)}
    transition_total = 0
    session_ids: set[str] = set()

    for directory in directories:
        resolved = directory.resolve()
        report.dataset_folders.append(str(resolved))
        if not directory.is_dir():
            report.errors.append(f"Dataset folder does not exist: {directory}")
            continue
        frames_dir = directory / "frames"
        actions_path = directory / "actions.jsonl"
        if not frames_dir.is_dir():
            report.errors.append(f"{directory} is missing a frames folder.")
            continue
        if not actions_path.is_file():
            report.errors.append(f"{directory} is missing actions.jsonl.")
            continue
        frame_files = sorted(path for path in frames_dir.iterdir() if path.is_file() and path.suffix.casefold() in IMAGE_EXTENSIONS)
        report.frames += len(frame_files)
        if not frame_files:
            report.errors.append(f"{directory} has no image frames.")
        indexed_frames = [index for index in (_numeric_frame_index(path) for path in frame_files) if index is not None]
        frame_paths_by_index: dict[int, list[Path]] = {}
        for frame_file in frame_files:
            index = _numeric_frame_index(frame_file)
            if index is not None:
                frame_paths_by_index.setdefault(index, []).append(frame_file)
        if indexed_frames:
            gaps = [
                (left, right) for left, right in zip(indexed_frames, indexed_frames[1:])
                if right != left + 1
            ]
            if gaps:
                report.warnings.append(f"{directory} has frame ordering gaps such as {gaps[0][0]} to {gaps[0][1]}; invalid transitions will be skipped.")
        rows_by_session: dict[str, list[dict[str, Any]]] = {}
        seen_keys: set[tuple[str, int]] = set()
        for line_number, line in enumerate(actions_path.read_text(encoding="utf-8").splitlines(), 1):
            line = line.strip()
            if not line:
                continue
            report.metadata_rows += 1
            try:
                row = json.loads(line)
            except json.JSONDecodeError:
                report.errors.append(f"{actions_path.name} line {line_number} is not valid JSON.")
                continue
            filename = str(row.get("filename", "")).strip()
            if not filename:
                report.errors.append(f"{actions_path.name} line {line_number} has no frame filename.")
                continue
            frame_path = frames_dir / filename
            if not frame_path.is_file():
                try:
                    frame_index = int(row.get("frame_index"))
                except (TypeError, ValueError):
                    report.warnings.append(
                        f"{actions_path.name} line {line_number} points to missing frame {filename}; skipping row."
                    )
                    continue
                candidates = frame_paths_by_index.get(frame_index, [])
                if len(candidates) == 1:
                    frame_path = candidates[0]
                    report.warnings.append(
                        f"{actions_path.name} line {line_number} uses {frame_path.name} for missing legacy filename {filename}."
                    )
                else:
                    report.warnings.append(
                        f"{actions_path.name} line {line_number} points to missing frame {filename}; skipping row."
                    )
                    continue
            if verify_images:
                try:
                    with Image.open(frame_path) as image:
                        image.verify()
                    with Image.open(frame_path) as image:
                        size = image.size
                except Exception as exc:
                    report.errors.append(f"Broken image file {frame_path.name}: {exc}")
                    continue
                if seen_resolution is None:
                    seen_resolution = size
                    report.resolution = f"{size[0]}x{size[1]}"
                elif size != seen_resolution:
                    report.errors.append(f"Inconsistent frame resolution: {frame_path.name} is {size[0]}x{size[1]}, expected {seen_resolution[0]}x{seen_resolution[1]}.")
            unexpected = sorted(set(row) - SUPPORTED_ACTIONS - METADATA_FIELDS)
            if unexpected:
                report.errors.append(f"{actions_path.name} line {line_number} contains unsupported action field(s): {', '.join(unexpected[:6])}.")
            missing = sorted(name for name in REQUIRED_CANONICAL_ACTIONS if name not in row)
            if missing:
                report.errors.append(f"{actions_path.name} line {line_number} is missing action label(s): {', '.join(missing)}.")
                continue
            try:
                frame_index = int(row.get("frame_index"))
            except (TypeError, ValueError):
                report.errors.append(f"{actions_path.name} line {line_number} has an invalid frame_index.")
                continue
            session_id = str(row.get("session_id") or f"legacy-{directory.name}").strip()
            key = (session_id, frame_index)
            if key in seen_keys:
                report.errors.append(f"{actions_path.name} repeats frame_index {frame_index} in session {session_id}.")
                continue
            seen_keys.add(key)
            action_active = False
            for name in action_counts:
                try:
                    value = float(row.get(name, 0))
                except (TypeError, ValueError):
                    report.errors.append(f"{actions_path.name} line {line_number} has invalid {name} action value.")
                    value = 0.0
                if abs(value) > (0.5 if name in BINARY_ACTIONS else 0.02):
                    action_counts[name] += 1
                    action_active = True
            if not action_active:
                report.idle_rows += 1
            row["_session_id"] = session_id
            rows_by_session.setdefault(session_id, []).append(row)
            session_ids.add(f"{resolved}:{session_id}")
            report.valid_rows += 1
        for session_id, rows in rows_by_session.items():
            if not rows:
                report.errors.append(f"{directory} has an empty sequence {session_id}.")
                continue
            rows.sort(key=lambda item: int(item["frame_index"]))
            for left, right in zip(rows, rows[1:]):
                if int(right["frame_index"]) != int(left["frame_index"]) + 1:
                    report.warnings.append(f"{directory} session {session_id} has an ordering gap at frame {left['frame_index']}; invalid transitions will be skipped.")
            transition_total += sum(
                1 for left, right in zip(rows, rows[frame_gap:])
                if int(right["frame_index"]) == int(left["frame_index"]) + frame_gap
            )

    report.sessions = len(session_ids)
    report.action_counts = action_counts
    report.valid_transitions = transition_total
    if report.metadata_rows != report.frames:
        report.warnings.append(f"Frame and label counts differ: {report.frames} frame files, {report.metadata_rows} action rows.")
    if report.valid_rows < 2:
        report.errors.append("The dataset needs at least two valid labelled frames.")
    if report.valid_transitions < 1:
        report.errors.append(f"No valid frame transitions were found for prediction horizon {frame_gap}.")
    if report.valid_rows and not any(action_counts.values()):
        report.errors.append("No non-idle action labels were found. Record idle plus at least one active control.")
    return report


def inspect_oasis_dataset(
    value: str | list[str] | tuple[str, ...], *, frame_gap: int = 1,
) -> OasisDatasetReport:
    """Quickly inspect training-relevant labels and pacing without decoding images."""
    return validate_oasis_dataset(value, frame_gap=frame_gap, verify_images=False)