EWMBench/evaluate.py at main · AgibotTech/EWMBench · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
import torch
import os
from EWMBench import EmbodiedWorldModelBenchmark
from EWMBench.distributed import dist_init, print0
from datetime import datetime
import argparse
import json
import yaml

def parse_args():

    CUR_DIR = os.path.dirname(os.path.abspath(__file__))
    parser = argparse.ArgumentParser(description='EmbodiedWorldModelBenchmark', formatter_class=argparse.RawTextHelpFormatter)

    parser.add_argument(
        "--dimension",
        nargs='+',
        required=True,
        help="list of evaluation dimensions, usage: --dimension <dim_1> <dim_2>",
    )

    parser.add_argument(
        "--config_path",
        type=str, default="",
        help="Path to config JSON that defines per-dimension settings"
        )

    parser.add_argument(
        "--overwrite",
        action="store_true"
    )

    args = parser.parse_args()
    return args

def load_config(config_path):
    with open(config_path, 'r') as f:
        config = yaml.safe_load(f)
    return config


def main():
    args = parse_args()

    dist_init()
    print0(f'args: {args}')
    device = torch.device("cuda")


    config = load_config(args.config_path)

    #base_path
    save_path = config['save_path']
    data_base = config['data']['val_base']
    gt_path = config['data']['gt_path']
    # data_name = config['model_name']
    data_name = os.path.basename(data_base).replace("_dataset", "")


    kwargs = {}

    #dimension
    for dim in args.dimension:
        if dim == 'semantics':
            kwargs[f"{dim}_caption_model_ckpt"] = config.get('ckpt',{}).get(dim, {}).get('caption', None)
            kwargs[f"{dim}_clip_model_ckpt"] = config.get('ckpt',{}).get(dim, {}).get('CLIP', None)
            kwargs[f"{dim}_bleus_model_ckpt"] = config.get(dim, {}).get('BLEUs', None)

            key_caption = f"{dim}_caption_model_ckpt"
            key_clip = f"{dim}_clip_model_ckpt"
            key_bleus = f"{dim}_bleus_model_ckpt"

            print(f"{dim}: caption = {kwargs[key_caption]}, CLIP = {kwargs[key_clip]}, BLEUs = {kwargs[key_bleus]}")
        elif dim in ['psnr', 'ssim']:
            pass

        else:
            kwargs[f"{dim}_model_ckpt"] = config.get('ckpt',{}).get(dim, None)

            key_model_ckpt = f"{dim}_model_ckpt"
            print(f"{dim}: ckpt = {kwargs[key_model_ckpt]}")

    print0(f'start evaluation')
    my_benchmark = EmbodiedWorldModelBenchmark(device, save_path)

    my_benchmark.evaluate(
        data_base = data_base,
        data_name = data_name,
        dimension_list = args.dimension,
        local=True,
        gt_path=gt_path,
        overwrite=args.overwrite,
        **kwargs
    )
    print0('done')


if __name__ == "__main__":
    main()