evaluator.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the CC BY-NC 4.0 license found in the
# LICENSE file in the root directory of this source tree.

import os
import tqdm
import logging
from typing import Dict, Any, Tuple, List

from data.load_frames import get_frames_from_video
from data.video_roots import VIDEO_ROOTS

logger = logging.getLogger(__name__)


class WorldPredictionEvaluator:
    """
    A unified evaluator supporting:
        (1) Single-step World Modeling (WM)
        (2) Multi-step Procedural Planning (PP)

    It loops over each sample in the provided annotation data and:
        - Constructs absolute video paths based on dataset_name + a VIDEO_ROOTS dictionary.
        - Loads the initial and final state frames.
        - Depending on the chosen task:
            * WM: Prepares candidate_actions, calls model.select_action(...)
            * PP: Prepares candidate_plans & candidate_actions, calls model.select_plan(...)
        - Compares model's prediction with ground_truth, tracks accuracy.

    Example Usage:
    -------------
    >>> import json
    >>> from models.example import ExampleModel
    >>> evaluator = WorldPredictionEvaluator(annotation_data, task="WM")
    >>> model = ExampleModel()
    >>> accuracy, results = evaluator.evaluate(model)

    Notes:
    ------
    - If you want to slice the dataset or do parallel processing, you can adapt
      this approach to your needs.
    - We assume all annotation-file entries match the structures described below.
    """

    def __init__(
        self,
        data: Dict[str, Dict[str, Any]],
        task: str = "WM",
        start_index: int = 0,
        end_index: int = None
    ):
        """
        Args:
            data (dict): Annotation data loaded from JSON.
            task (str): Either "WM" (World Modeling) or "PP" (Procedural Planning).
            start_index (int): Start index of samples to evaluate.
            end_index (int): End index (exclusive) of samples to evaluate.
        """
        if task not in ("WM", "PP"):
            raise ValueError("task must be 'WM' or 'PP'.")

        self.data = data
        self.task = task
        self.start_index = start_index
        self.end_index = end_index

        logger.info(f"WorldPredictionEvaluator initialized for WorldPrediction-{self.task}. "
                    f"Found {len(self.data)} datasets in annotation data.")

    def evaluate(self, model) -> Tuple[float, List[dict]]:
        """
        Perform evaluation over all samples in self.data using the provided model.

        The model must implement:
            - select_action(states, candidate_actions) -> str (for WM), or
            - select_plan(states, candidate_plans, candidate_actions) -> int (for PP).

        Returns:
            Tuple:
                accuracy (float): Fraction of correct predictions.
                results (List[dict]): One entry per evaluated sample.
        """
        logger.info(f"Starting evaluation for task={self.task}.")
        results = []
        total_count = 0
        correct_count = 0
        global_sample_index = 0

        for dataset_name, samples_dict in self.data.items():
            dataset_root = VIDEO_ROOTS.get(dataset_name)
            if dataset_root is None:
                logger.warning(f"No VIDEO_ROOT defined for dataset '{dataset_name}'. Skipping its samples.")
                continue

            sample_items = list(samples_dict.items())
            logger.info(f"Evaluating {len(sample_items)} samples in dataset '{dataset_name}'.")

            for sample_uid, sample_info in tqdm.tqdm(sample_items):
                if global_sample_index < self.start_index:
                    global_sample_index += 1
                    continue
                if self.end_index is not None and global_sample_index >= self.end_index:
                    break

                logger.info(f"[{global_sample_index}] Processing sample_uid='{sample_uid}' from dataset='{dataset_name}'.")
                global_sample_index += 1
                total_count += 1

                abs_video_path = os.path.join(dataset_root, sample_info["states"]["video"])
                init_img = get_frames_from_video(
                    video_path=abs_video_path,
                    timestamps=sample_info["states"]["segment_start_time"]
                )
                final_img = get_frames_from_video(
                    video_path=abs_video_path,
                    timestamps=sample_info["states"]["segment_end_time"]
                )
                states = {
                    "initial_state": init_img,
                    "final_state": final_img,
                    "segment_uid": sample_info["states"]["segment_uid"],
                    "sample_uid": sample_uid,
                }

                if self.task == "WM":
                    ground_truth = sample_info["ground_truth"]
                    candidate_actions = []
                    for cand in sample_info["candidates"]:
                        abs_cand_path = os.path.join(dataset_root, cand["video"])
                        candidate_actions.append({
                            "video": abs_cand_path,
                            "start_time": cand["segment_start_time"],
                            "end_time": cand["segment_end_time"],
                            "segment_uid": cand["segment_uid"]
                        })
                    prediction, info = model.select_action(states, candidate_actions)
                    is_correct = (prediction == ground_truth)
                else:
                    ground_truth = sample_info["ground_truth"]
                    candidate_plans = sample_info["candidates"]
                    seg_dict = sample_info["action_segments"]
                    candidate_actions = {}
                    # for seg_uid, seg_info in seg_dict.items():
                    for seg_uid in candidate_plans[0]: # arrange the candidate actions by option A
                        seg_info = seg_dict[seg_uid]
                        abs_seg_path = os.path.join(dataset_root, seg_info["video"])
                        candidate_actions[seg_uid] = {
                            "video": abs_seg_path,
                            "start_time": seg_info["segment_start_time"],
                            "end_time": seg_info["segment_end_time"],
                            "segment_uid": seg_uid
                        }
                    prediction, info = model.select_plan(states, candidate_plans, candidate_actions)
                    is_correct = (prediction == ground_truth)

                if is_correct:
                    correct_count += 1

                results.append({
                    "dataset_name": dataset_name,
                    "sample_uid": sample_uid,
                    "ground_truth": ground_truth,
                    "prediction": prediction,
                    "is_correct": is_correct,
                    "info": info,
                })

                current_accuracy = correct_count / total_count if total_count > 0 else 0.0
                logger.info(f"Sample '{sample_uid}': is_correct={is_correct}, "
                            f"Accuracy so far: {current_accuracy:.3f}")

        accuracy = correct_count / total_count if total_count > 0 else 0.0
        logger.info(f"Evaluation complete. Overall accuracy={accuracy:.3f} "
                    f"({correct_count}/{total_count} correct).")
        return accuracy, results