{
  "states": {
    "segment_uid": "segment|COIN|U9bZX2wABWQ_000",
    "video": "U9bZX2wABWQ.mp4",
    "segment_start_time": 23.7,
    "segment_end_time": 42.3
  },
  "ground_truth": "segment|COIN|pCCVwljBs7M_005",
  "candidates": [
    {
      "video": "pCCVwljBs7M.mp4",
      "segment_start_time": 112.2,
      "segment_end_time": 123.8,
      "action_label": "set up the brackets",
      "segment_uid": "segment|COIN|pCCVwljBs7M_005"
    },
    ...
  ]
}

{
  "states": {
    "segment_uid": "PP|segment|COIN|dBukB_1cXBE_003|-|segment|COIN|dBukB_1cXBE_005",
    "video": "dBukB_1cXBE.mp4",
    "segment_start_time": 60.7,
    "segment_end_time": 76.3
  },
  "ground_truth": 3,
  "candidates": [
    ["segment|COIN|AVbdqzpyAxk_007", "segment|COIN|AVbdqzpyAxk_004", "segment|COIN|dMSqm5BB5jA_000"],
    ["segment|COIN|AVbdqzpyAxk_004", "segment|COIN|AVbdqzpyAxk_007", "segment|COIN|dMSqm5BB5jA_000"],
    ...
  ],
  "action_segments": {
    "segment|COIN|dMSqm5BB5jA_000": {
      "video": "dMSqm5BB5jA.mp4",
      "segment_start_time": 28.7,
      "segment_end_time": 42.3,
      "action_label": "take out the shell"
    },
    ...
  }
}

VIDEO_ROOTS = {
    "COIN": "/path/to/COIN/",
    "CrossTask": "/path/to/CrossTask/",
    "EgoExo4D": "/path/to/EgoExo4D/",
    "IKEAASM": "/path/to/IKEA-ASM/",
    "EPIC-KITCHENS-100": "/path/to/EPIC-KITCHENS-100/"
}

import json
from evaluator import WorldPredictionEvaluator
from model import VLM

# Load annotations
data = json.load(open("WorldPrediction-WM.json", "r"))

# Initialize evaluator
evaluator = WorldPredictionEvaluator(
    data=data,
    task="WM",
    start_index=0,
    end_index=10  # Evaluate first 10 samples
)

# Load model
config = json.load(open("configs/vlm/qwenvl/Qwen2.5-VL-3B-Instruct-4f.json", "r"))
model = VLM(**config)

# Run evaluation
accuracy, results = evaluator.evaluate(model)
print(f"Accuracy: {accuracy:.3f}")

class MyCustomModel:
    def __init__(self, **kwargs):
        """Initialize your model with configuration parameters."""
        pass

    def select_action(self, states: dict, candidate_actions: list) -> tuple[str, dict]:
        """
        Single-step action selection for World Modeling.

        Args:
            states: {
                "initial_state": PIL.Image,  # Initial world state
                "final_state": PIL.Image,    # Final world state
                "segment_uid": str,          # Unique identifier
                "sample_uid": str,           # Sample identifier
            }
            candidate_actions: List of dicts, each containing:
                - "video": str (path to video file)
                - "start_time": float
                - "end_time": float
                - "segment_uid": str

        Returns:
            Tuple of (selected_segment_uid: str, info: dict)
        """
        pass

    def select_plan(self, states: dict, candidate_plans: list,
                    candidate_actions: dict) -> tuple[int, dict]:
        """
        Multi-step plan selection for Procedural Planning.

        Args:
            states: Same as select_action
            candidate_plans: List of plans, where each plan is a list of segment_uids
            candidate_actions: Dict mapping segment_uid to action info

        Returns:
            Tuple of (selected_plan_index: int, info: dict)
        """
        pass

from model.my_custom_model import MyCustomModel