Function main

evaluation.py:209–328

()

            self.handleError(record)

def main():
    import argparse
    
    parser = argparse.ArgumentParser(description="Evaluate World Model Benchmark")
    parser.add_argument("--judge", type=str, required=True, help="Path to judge model checkpoint")
    parser.add_argument("--video_dir", type=str, required=True, help="Path to generated video directory")
    parser.add_argument("--model_name", type=str, required=True, help="Tested model name")
    parser.add_argument("--save_name", type=str, default="worldmodelbench_results", help="Path to save evaluation results")
    parser.add_argument("--cot", action="store_true", help="Enable Chain-of-Thought output")
    parser.add_argument("--no-save", action="store_true", help="Disable saving results")
    
    args = parser.parse_args()
    
    # Setup logging with custom Rich handler
    logging.basicConfig(
        level=logging.INFO,
        format="%(message)s",
        handlers=[RichLogHandler()]
    )
    logger = logging.getLogger(__name__)

    # Initialize evaluator
    config = EvaluationConfig()
    evaluator = WorldModelEvaluator(args.judge, args.video_dir, config)
    printer = ResultsPrinter()
    
    # Load validation set with status message
    printer.console.print("[bold]Loading validation set...[/bold]")
    validation_set = load("./worldmodelbench.json")
    
    # Check for existing results
    save_path = f"{args.save_name}_cot.json" if args.cot else f"{args.save_name}.json"
    if os.path.exists(save_path):
        printer.console.print("[bold yellow]Loading existing results...[/bold yellow]")
        results = load(save_path)
        try:
            preds, accs = results["preds"], results["accs"]
        except KeyError:
            raise KeyError("Expected keys not found in results file")
    else:
        printer.console.print("[bold green]Starting new evaluation...[/bold green]")
        preds = {}
        accs = defaultdict(list)
        
        # Create a single progress instance for all operations
        with Progress(
            "[progress.description]{task.description}",
            BarColumn(),
            "[progress.percentage]{task.percentage:>3.0f}%",
            TimeRemainingColumn(),
            console=printer.console
        ) as progress:
            # Main task for video processing
            video_task = progress.add_task("Processing videos", total=len(validation_set))

            for vid, v_i in tqdm(enumerate(validation_set), total=len(validation_set)):
                video_name = Path(v_i["first_frame"]).stem
                video = evaluator._load_video(video_name)