dataset.py

import json
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional


@dataclass
class CustomEpisode:
    episode_id: str
    scene_id: str
    start_position: List[float]
    start_rotation: List[float]
    target_object_ids: List[int]
    target_mode: str = "any"
    target_count: Optional[int] = None
    shortest_path_distance: float = 0.0
    target_positions: Dict[int, List[float]] = field(default_factory=dict)
    natural_language: List[str] = field(default_factory=list)
    query_program: Dict[str, Any] = field(default_factory=dict)
    goal_type: str = "description"
    goal_text: str = ""
    object_category: Optional[str] = None
    extras: Dict[str, Any] = field(default_factory=dict)


class CustomDataset:
    """Lightweight wrapper for custom multi-object navigation episodes."""

    def __init__(self, episodes: List[CustomEpisode]):
        self.episodes = episodes

    def __len__(self) -> int:
        return len(self.episodes)

    def __getitem__(self, idx: int) -> CustomEpisode:
        return self.episodes[idx]

    def get_episode(self, idx: int) -> CustomEpisode:
        return self.episodes[idx]

    def __iter__(self):
        return iter(self.episodes)


def _load_json(path: str) -> Any:
    path_obj = Path(path)
    with path_obj.open("r", encoding="utf-8") as file:
        return json.load(file)


def _normalize_list(value: Any) -> List[Any]:
    if value is None:
        return []
    if isinstance(value, list):
        return value
    return [value]


def _load_goal_metadata(goals_path: Optional[str]) -> tuple[Dict[int, List[float]], set[str]]:
    if goals_path is None:
        return {}, set()

    goals_json = _load_json(goals_path)
    goals = goals_json.get("goals", {}) if isinstance(goals_json, dict) else {}
    positions: Dict[int, List[float]] = {}
    categories: set[str] = set()
    for goal in goals.values():
        object_id = goal.get("object_id")
        position = goal.get("position")
        category = goal.get("category")
        if category is not None and str(category).strip():
            categories.add(str(category).strip())
        if object_id is None or not isinstance(position, list) or len(position) < 3:
            continue
        positions[int(object_id)] = [float(x) for x in position[:3]]
    return positions, categories


def _simple_category_value(query_program: Dict[str, Any]) -> Optional[str]:
    """Return the object category when a query is exactly category == value."""
    if not isinstance(query_program, dict):
        return None
    allowed_program_keys = {"select", "where", "target_count"}
    if any(key not in allowed_program_keys for key in query_program.keys()):
        return None
    conditions = query_program.get("where")
    if not isinstance(conditions, list) or len(conditions) != 1:
        return None
    condition = conditions[0]
    if not isinstance(condition, dict):
        return None
    if condition.get("field") != "category" or condition.get("op") != "==":
        return None
    value = condition.get("value")
    if value is None:
        return None
    value = str(value).strip()
    return value or None


def load_custom_dataset(
    episodes_path: str,
    queries_path: Optional[str] = None,
    goals_path: Optional[str] = None,
    goal_type: str = "description",
) -> CustomDataset:
    """Load episodes and attach natural language or query metadata when available."""
    goal_type = str(goal_type or "description").lower()
    if goal_type not in {"description", "object"}:
        raise ValueError(f"Unknown goal_type: {goal_type}")

    episodes_json = _load_json(episodes_path)
    if isinstance(episodes_json, dict) and "episodes" in episodes_json:
        episodes_json = episodes_json["episodes"]

    goal_positions, category_vocab = _load_goal_metadata(goals_path)

    query_map: Dict[str, Dict[str, Any]] = {}
    queries_list: List[Dict[str, Any]] = []
    if queries_path is not None:
        queries_json = _load_json(queries_path)
        if isinstance(queries_json, dict):
            if "queries" in queries_json:
                queries_json = queries_json["queries"]
            elif "episodes" in queries_json:
                queries_json = queries_json["episodes"]
        queries_list = list(queries_json) if isinstance(queries_json, list) else []
        for query in queries_list:
            query_id = query.get("query_id") or query.get("episode_id")
            if query_id is not None:
                query_map[str(query_id)] = query

    episodes: List[CustomEpisode] = []
    for index, episode_data in enumerate(episodes_json):
        episode_id = str(episode_data.get("episode_id", index))
        query_id = episode_data.get("query_id")
        query_data = query_map.get(str(query_id)) if query_id is not None else None
        if query_data is None:
            query_data = query_map.get(episode_id)
        if query_data is None and len(queries_list) == len(episodes_json):
            query_data = queries_list[index]

        natural_language = []
        if query_data is not None:
            raw_nl = query_data.get("natural_language")
            natural_language = _normalize_list(raw_nl)
        if not natural_language:
            natural_language = _normalize_list(episode_data.get("instruction"))

        query_program = query_data.get("query_program", {}) if query_data else {}
        object_category = _simple_category_value(query_program)
        is_object_nav_episode = (
            object_category is not None
            and (not category_vocab or object_category in category_vocab)
        )
        if goal_type == "object" and not is_object_nav_episode:
            continue

        target_object_ids = episode_data.get("target_object_ids")
        if not target_object_ids and query_data is not None:
            target_object_ids = query_data.get("target_object_ids")
        target_object_ids = _normalize_list(target_object_ids)
        target_object_ids = [int(x) for x in target_object_ids if x is not None]

        target_mode = str(episode_data.get("target_mode", "any") or "any").lower()
        target_count = episode_data.get("target_count")
        if target_count is None and query_data is not None:
            target_count = query_data.get("target_count")
        if target_count is not None:
            target_count = int(target_count)
        shortest_path_distance = float(episode_data.get("shortest_path_distance", 0.0) or 0.0)

        episode = CustomEpisode(
            episode_id=episode_id,
            scene_id=str(episode_data.get("scene_id", "")),
            start_position=[float(x) for x in episode_data.get("start_position", [])],
            start_rotation=[float(x) for x in episode_data.get("start_rotation", [])],
            target_object_ids=target_object_ids,
            target_mode=target_mode,
            target_count=target_count,
            shortest_path_distance=shortest_path_distance,
            target_positions={
                object_id: goal_positions[object_id]
                for object_id in target_object_ids
                if object_id in goal_positions
            },
            natural_language=natural_language,
            query_program=query_program,
            goal_type=goal_type,
            goal_text=object_category if goal_type == "object" and object_category else "",
            object_category=object_category,
            extras={k: v for k, v in episode_data.items() if k not in {
                "episode_id",
                "scene_id",
                "start_position",
                "start_rotation",
                "target_object_ids",
                "target_mode",
                "target_count",
                "shortest_path_distance",
            }} | {
                "query_id": query_id,
                "object_nav_eligible": is_object_nav_episode,
                "category_vocab_size": len(category_vocab),
            },
        )
        episodes.append(episode)

    return CustomDataset(episodes)