mdl2fluid.py 2.8 KB
Newer Older
xiebaiyuan's avatar
xiebaiyuan 已提交
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 101 102 103 104 105 106 107 108 109 110 111 112
import json
import framework_pb2 as framework_pb2
import op_types as types


def load_mdl(mdl_json_path):
    # print('mdl json path : ' + mdl_json_path)
    with open(mdl_json_path, 'r') as f:
        return json.load(f)


class Converter:
    'convert mdlmodel to fluidmodel'

    def __init__(self, mdl_json_path):
        self.mdl_json_path = mdl_json_path
        self.mdl_json = load_mdl(self.mdl_json_path)
        self.program_desc = framework_pb2.ProgramDesc()

        # print(json_dick)
        # layers = (json_dick['layer'])
        # for layer in layers:
        #     print(layer)

    def convert(self):
        print 'convert begin.....'
        # add block_desc
        block_desc = self.program_desc.blocks.add()
        block_desc.idx = 0
        block_desc.parent_idx = -1
        self.package_ops(block_desc)
        print 'blocks: '
        print self.program_desc.blocks

    def package_ops(self, block_desc):
        # add ops with layer
        if 'layer' in self.mdl_json:

            layers_ = self.mdl_json['layer']
            for layer in layers_:
                desc_ops = block_desc.ops.add()

                # print layer
                # for i in layer:
                #     print i
                if 'name' in layer:
                    l_name = layer['name']

                if 'weight' in layer:
                    l_weights = layer['weight']

                if 'param' in layer:
                    l_params = layer['param']

                if 'output' in layer:
                    l_outputs = layer['output']

                if 'input' in layer:
                    l_inputs = layer['input']
                    inputs_add = desc_ops.inputs.add()
                    for i in l_inputs:
                        # print i
                        inputs_add.parameter = ''
                        inputs_add.arguments.append(i)

                if 'type' in layer:
                    l_type = layer['type']
                    # print l_type
                    # print mdl2fluid_op_layer_dict.get(l_type)
                    desc_ops.type = types.mdl2fluid_op_layer_dict.get(l_type)


mdl_path = "multiobjects/YOLO_Universal.json"
# print mdl_path
# # model
# mdl_model = load_mdl(mdl_path)
# for key in mdl_model:
#     print key
#
# # layer
# layers = mdl_model['layer']
# print layers
#
# for layer in layers:
#     print layer
#     for i in layer:
#         print i
#     if 'name' in layer:
#         l_name = layer['name']
#
#     if 'weight' in layer:
#         l_weights = layer['weight']
#
#     if 'param' in layer:
#         l_params = layer['param']
#
#     if 'output' in layer:
#         l_outputs = layer['output']
#
#     if 'input' in layer:
#         l_inputs = layer['input']
#
#     if 'type' in layer:
#         l_type = layer['type']
#
# print mdl_model['matrix']
#
# package()


converter = Converter(mdl_path)
converter.convert()