1import sys as _sys
2import google.protobuf.text_format as text_format
3from six import text_type as _text_type
4
5
6
7def _convert(args):
8    if args.framework == 'caffe':
9        raise NotImplementedError("Destination [Caffe] is not implemented yet.")
10
11    elif args.framework == 'keras':
12        raise NotImplementedError("Destination [Keras] is not implemented yet.")
13
14    elif args.framework == 'tensorflow':
15        raise NotImplementedError("Destination [Tensorflow] is not implemented yet.")
16
17    elif args.framework == 'cntk':
18        raise NotImplementedError("Destination [CNTK] is not implemented yet.")
19
20    elif args.framework == 'coreml':
21        from mmdnn.conversion.coreml.coreml_emitter import CoreMLEmitter
22        assert args.inputNetwork is not None
23        assert args.inputWeight is not None
24        emitter = CoreMLEmitter(args.inputNetwork, args.inputWeight)
25        model, in_, out_ = emitter.gen_model(
26            args.inputNames,
27            args.outputNames,
28            image_input_names = set(args.imageInputNames) if args.imageInputNames else None,
29            is_bgr = args.isBGR,
30            red_bias = args.redBias,
31            blue_bias = args.blueBias,
32            green_bias = args.greenBias,
33            gray_bias = args.grayBias,
34            image_scale = args.scale,
35            class_labels = args.classInputPath if args.classInputPath else None,
36            predicted_feature_name = args.predictedFeatureName)
37
38        """
39        from google.protobuf import text_format
40        with open(args.output+'.txt', 'w') as f:
41            f.write(text_format.MessageToString(model))
42        """
43
44        with open(args.output, 'wb') as f:
45            model = model.SerializeToString()
46            f.write(model)
47
48
49        return 0
50
51    elif args.framework == 'pytorch':
52        if not args.dstWeightPath or not args.IRWeightPath:
53            raise ValueError("Need to set a target weight filename.")
54        from mmdnn.conversion.pytorch.pytorch_emitter import PytorchEmitter
55        emitter = PytorchEmitter((args.IRModelPath, args.IRWeightPath))
56
57    elif args.framework == 'mxnet':
58        from mmdnn.conversion.mxnet.mxnet_emitter import MXNetEmitter
59        if args.IRWeightPath == None:
60            emitter = MXNetEmitter(args.IRModelPath)
61        else:
62            emitter = MXNetEmitter((args.IRModelPath, args.IRWeightPath, args.inputShape, args.dstWeightPath))
63
64    else:
65        assert False
66
67    emitter.run(args.output)
68
69    return 0
70
71
72def _get_parser():
73    import argparse
74
75    parser = argparse.ArgumentParser(description='Convert IR model file formats to other format.')
76
77    parser.add_argument(
78        '-f', '--framework', type=_text_type, choices=['coreml'], required=True,
79        help='Format of model at srcModelPath (default is to auto-detect).'
80    )
81
82    parser.add_argument(
83        '-in', '--inputNetwork',
84        type=_text_type,
85        required=True,
86        help='Path of the IR network architecture file.')
87
88    parser.add_argument(
89        '-iw', '--inputWeight',
90        type=_text_type,
91        required=True,
92        help='Path to the IR network weight file.')
93
94    parser.add_argument(
95        '-o', '--output',
96        type=_text_type,
97        required=True,
98        help='Path to save the destination model')
99
100    # Caffe
101    parser.add_argument(
102        '--phase', type=_text_type, choices=['train', 'test'], default='test',
103        help='[Caffe] Convert phase (train/test) for destination toolkits.'
104    )
105
106    # For CoreML
107    parser.add_argument('--inputNames', type=_text_type, nargs='*', help='Names of the feature (input) columns, in order (required for keras models).')
108    parser.add_argument('--outputNames', type=_text_type, nargs='*', help='Names of the target (output) columns, in order (required for keras models).')
109    parser.add_argument('--imageInputNames', type=_text_type, default=[], action='append', help='Label the named input as an image. Can be specified more than once for multiple image inputs.')
110    parser.add_argument('--isBGR', action='store_true', default=False, help='True if the image data in BGR order (RGB default)')
111    parser.add_argument('--redBias', type=float, default=0.0, help='Bias value to be added to the red channel (optional, default 0.0)')
112    parser.add_argument('--blueBias', type=float, default=0.0, help='Bias value to be added to the blue channel (optional, default 0.0)')
113    parser.add_argument('--greenBias', type=float, default=0.0, help='Bias value to be added to the green channel (optional, default 0.0)')
114    parser.add_argument('--grayBias', type=float, default=0.0, help='Bias value to be added to the gray channel for Grayscale images (optional, default 0.0)')
115    parser.add_argument('--scale', type=float, default=1.0, help='Value by which the image data must be scaled (optional, default 1.0)')
116    parser.add_argument('--classInputPath', type=_text_type, default='', help='Path to class labels (ordered new line separated) for treating the neural network as a classifier')
117    parser.add_argument('--predictedFeatureName', type=_text_type, default='class_output', help='Name of the output feature that captures the class name (for classifiers models).')
118    return parser
119
120
121def _main():
122    parser=_get_parser()
123    args = parser.parse_args()
124    ret = _convert(args)
125    _sys.exit(int(ret)) # cast to int or else the exit code is always 1
126
127
128if __name__ == '__main__':
129    _main()
130