Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
generate_md_table.py 2.71 KB
Newer Older
limm's avatar
add tools and service module  
limm committed
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 
# Copyright (c) OpenMMLab. All rights reserved.
import argparse
import os
import os.path as osp

import yaml
from mmengine import Config

from mmdeploy.utils import get_backend, get_task_type, load_config


def parse_args():
    parser = argparse.ArgumentParser(
        description='from yaml export markdown table')
    parser.add_argument('yml_file', help='input yml config path')
    parser.add_argument('output', help='output markdown file path')
    parser.add_argument(
        '--backends',
        nargs='+',
        help='backends you want to generate',
        default=[
            'onnxruntime', 'tensorrt', 'torchscript', 'pplnn', 'openvino',
            'ncnn'
        ])
    args = parser.parse_args()
    return args


def main():
    args = parse_args()
    assert osp.exists(args.yml_file), f'File not exists: {args.yml_file}'
    output_dir, _ = osp.split(args.output)
    if output_dir:
        os.makedirs(output_dir, exist_ok=True)
    header = ['model', 'task'] + args.backends
    aligner = [':--'] * 2 + [':--:'] * len(args.backends)

    def write_row_f(writer, row):
        writer.write('|' + '|'.join(row) + '|\n')

    print(f'Processing{args.yml_file}')
    with open(args.yml_file, 'r') as reader, open(args.output, 'w') as writer:
        config = yaml.load(reader, Loader=yaml.FullLoader)
        config = Config(config)
        write_row_f(writer, header)
        write_row_f(writer, aligner)
        repo_url = config.globals.repo_url
        for i in range(len(config.models)):
            name = config.models[i].name
            model_configs = config.models[i].model_configs
            pipelines = config.models[i].pipelines
            config_url = osp.join(repo_url, model_configs[0])
            config_url, _ = osp.split(config_url)
            support_backends = {b: 'N' for b in args.backends}
            deploy_config = [
                pipelines[i].deploy_config for i in range(len(pipelines))
            ]
            cfg = [
                load_config(deploy_config[i])
                for i in range(len(deploy_config))
            ]
            task = [
                get_task_type(cfg[i][0]).value
                for i in range(len(deploy_config))
            ]
            backend_type = [
                get_backend(cfg[i][0]).value
                for i in range(len(deploy_config))
            ]
            for i in range(len(deploy_config)):
                support_backends[backend_type[i]] = 'Y'
            support_backends = [support_backends[i] for i in args.backends]
            model_name = f'[{name}]({config_url})'
            row = [model_name, task[i]] + support_backends

            write_row_f(writer, row)
        print(f'Save to {args.output}')


if __name__ == '__main__':
    main()