ir_name_map_util.py

import os
import sys
import json
import xml.etree.ElementTree as ET

def get_json_name_map(model_path_before, model_path_after, output_path):
    attrs_to_check = ['type']
    attr_to_map = 'name'

    root_before = ET.parse(model_path_before).getroot()
    root_after = ET.parse(model_path_after).getroot()

    layers_before = root_before[0]
    layers_after = root_after[0]
    assert len(layers_before) == len(layers_after),\
            f'Amount of layer before ({len(layers_before)} is not match amount of layers after ({len(layers_after)})'

    result = {}
    for idx, (l_before, l_after) in enumerate(zip(layers_before, layers_after)):
        before_name = l_before.get(attr_to_map)
        after_name = l_after.get(attr_to_map)
        for attr in attrs_to_check:
            before_attr = l_before.get(attr)
            after_attr = l_after.get(attr)
            assert before_attr == after_attr, \
                    f'for layer {before_name} -> {after_name} attribute {attr} does not match:\n'\
                    f'{before_attr} vs {after_attr}'
        result[before_name] = after_name

    print(f'Saving results to {output_path}')
    with open(output_path, 'w') as f:
        f.write(json.dumps(result))
    print('Done!')


def map_names(path_to_names, path_to_map, output_path):
    with open(path_to_map, 'r') as f:
        name_map = json.loads('\n'.join(f.readlines()))
    with open(path_to_names) as f:
        names = f.read().splitlines()
    result = []
    for name in names:
        result.append(name_map[name])
    with open(output_path, 'w') as f:
        f.writelines(result)


if __name__ == '__main__':
    options = ['gen', 'map']
    if len(sys.argv) == 1 or sys.argv[1] not in options:
        print(f'options: {options}')
        exit()
    if sys.argv[1] == options[0]:
        get_json_name_map(*sys.argv[2:])
    if sys.argv[1] == options[1]:
        map_names(*sys.argv[2:])