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)